forked from Karylab-cklius/vllm
Compare commits
217
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5ea7cac55b | ||
|
|
be3476447f | ||
|
|
f1473092c4 | ||
|
|
4dcbae8670 | ||
|
|
65dac3a770 | ||
|
|
0ba2500ef0 | ||
|
|
ef576befd2 | ||
|
|
35e4a36107 | ||
|
|
da5803d46e | ||
|
|
75ddfaf909 | ||
|
|
f36fe52add | ||
|
|
391d918d4d | ||
|
|
fca040885b | ||
|
|
a7fd4c7482 | ||
|
|
5131691063 | ||
|
|
80e00e5ac6 | ||
|
|
ae10e855ab | ||
|
|
530ee36a0d | ||
|
|
d835ad572c | ||
|
|
47d0597ca2 | ||
|
|
818cf61e91 | ||
|
|
c01618fdc8 | ||
|
|
823eaf667d | ||
|
|
f1f1259692 | ||
|
|
df13b5aef5 | ||
|
|
4938d44a3b | ||
|
|
37bf988c2f | ||
|
|
9459fc6471 | ||
|
|
5245c80564 | ||
|
|
9bc266d923 | ||
|
|
5c9f6557d7 | ||
|
|
dcfebf93f4 | ||
|
|
752bd10647 | ||
|
|
2730b657c4 | ||
|
|
1dcbbd9cac | ||
|
|
ace9fda495 | ||
|
|
ef0aa7ca2f | ||
|
|
e6d1310b2a | ||
|
|
ac5f38a0f7 | ||
|
|
b6ff8a2f50 | ||
|
|
9243e0124e | ||
|
|
df362b2d6d | ||
|
|
7c2acd38b7 | ||
|
|
a287eb163f | ||
|
|
e94243893d | ||
|
|
29c0ec4d63 | ||
|
|
c7ce03bcbd | ||
|
|
c233d90aa8 | ||
|
|
d96aee0951 | ||
|
|
c71a583aa9 | ||
|
|
f12b80c6ef | ||
|
|
da64db78b9 | ||
|
|
425c4eafb0 | ||
|
|
02c01f442b | ||
|
|
fae543015c | ||
|
|
c9be3a8aa1 | ||
|
|
41ea2dd44a | ||
|
|
088c0be268 | ||
|
|
fcd2255d16 | ||
|
|
b5433b6f50 | ||
|
|
cc25f028b7 | ||
|
|
c4cd2bd544 | ||
|
|
5784507da4 | ||
|
|
bf578e1abd | ||
|
|
efed8a1e83 | ||
|
|
11d291511a | ||
|
|
877dae9c68 | ||
|
|
c4dd6d78fd | ||
|
|
ce2aecc4dc | ||
|
|
f38f3d11fb | ||
|
|
d4b4562917 | ||
|
|
7b3192523e | ||
|
|
4c6e2e4b30 | ||
|
|
8502958810 | ||
|
|
ce4bdcbda4 | ||
|
|
d5b1ec2684 | ||
|
|
867ff69733 | ||
|
|
109b736b86 | ||
|
|
69d4f5ef63 | ||
|
|
426d48bfa1 | ||
|
|
26c909ed74 | ||
|
|
fb1d8ccaf5 | ||
|
|
9354f22204 | ||
|
|
17fdd42100 | ||
|
|
472d330c21 | ||
|
|
3b6c96a101 | ||
|
|
4d4e04f452 | ||
|
|
67fe73b2b4 | ||
|
+1 |
ee8f36d0b3 | ||
|
+1 |
f3e9497e92 | ||
|
|
fe784ff22e | ||
|
|
b88abb5036 | ||
|
|
67f9046e4a | ||
|
|
f17be06fbe | ||
|
|
2cab53ddee | ||
|
|
ab0a20d151 | ||
|
|
4a394bfcda | ||
|
|
c95c663049 | ||
|
|
ab3c1aedf3 | ||
|
+1 |
fb5ec0dc9e | ||
|
|
971dac2caa | ||
|
|
efa2e424f6 | ||
|
|
02bf9c7907 | ||
|
|
f61163e6c7 | ||
|
|
626c90b2d5 | ||
|
+1 |
251f7e478e | ||
|
|
ce65385618 | ||
|
|
7d56fe2adc | ||
|
|
75bdad40b5 | ||
|
|
d08eebad16 | ||
|
|
3e90d015ba | ||
|
|
7cd1d57b74 | ||
|
|
b8168e33e0 | ||
|
|
d803b44dbe | ||
|
|
530852f959 | ||
|
|
a317bc5739 | ||
|
|
a9531edfa6 | ||
|
|
8c3393f373 | ||
|
|
9f8cbfd8eb | ||
|
|
ea1d65fe6d | ||
|
|
f44f3d6f79 | ||
|
|
cc706b05a5 | ||
|
|
85e296950c | ||
|
|
dc9f845ddc | ||
|
|
12f2c515a7 | ||
|
+1 |
6570c9800c | ||
|
|
8bfd683901 | ||
|
|
7dc2698632 | ||
|
|
59b964f37d | ||
|
|
6a9f24aa8c | ||
|
|
ba47bb5be1 | ||
|
|
df8a0900df | ||
|
|
2db39c7049 | ||
|
|
3935829f89 | ||
|
|
7746961277 | ||
|
|
5de1add806 | ||
|
|
915dffaa5f | ||
|
|
81e13a0591 | ||
|
|
f95e3f0edb | ||
|
|
5a65ba5f17 | ||
|
|
9d1c695be5 | ||
|
|
3c1bc1fc0d | ||
|
|
3a5e88e629 | ||
|
|
0becb7486b | ||
|
|
2dab187f75 | ||
|
|
015b0320de | ||
|
|
4238b011a7 | ||
|
|
eb33ff34dd | ||
|
|
2bd8957627 | ||
|
|
3034c8d389 | ||
|
|
ecf4aa5ce2 | ||
|
|
49e777cf08 | ||
|
|
b7950e798f | ||
|
|
de100ffb62 | ||
|
|
43cd340247 | ||
|
|
1d99f0f421 | ||
|
|
0885b51981 | ||
|
|
6036bf110a | ||
|
|
2fa63e0fff | ||
|
|
61141ed265 | ||
|
|
05eed72aec | ||
|
|
5810e884f1 | ||
|
|
615834ee58 | ||
|
|
5811ed6a05 | ||
|
|
1b30ae4ca4 | ||
|
|
4e04bcbce6 | ||
|
|
66b6c684ab | ||
|
|
c0302d9497 | ||
|
|
313fae3e89 | ||
|
|
7aab6e2684 | ||
|
|
9dd2e72828 | ||
|
|
d119beb1b9 | ||
|
|
12a8057bfe | ||
|
|
e281ac663a | ||
|
|
adce068118 | ||
|
|
b6770d7b54 | ||
|
|
3b39fd284a | ||
|
|
6472131298 | ||
|
|
37aa52821d | ||
|
|
96d2ceda4b | ||
|
|
fdf2cf66d3 | ||
|
|
9b2be4e9a5 | ||
|
|
3ad85e0de4 | ||
|
|
4f7fffb92f | ||
|
|
6e073440b1 | ||
|
|
f7aadae5e5 | ||
|
|
442c421e79 | ||
|
|
0bd6b85a1f | ||
|
|
3ca242d1b6 | ||
|
|
7e950521b3 | ||
|
|
0f0f28b537 | ||
|
|
520a20ba4e | ||
|
|
9182e86971 | ||
|
|
313d01f507 | ||
|
|
05d4f8bba3 | ||
|
|
0b54201a04 | ||
|
|
32e632dfeb | ||
|
|
7ffb98e248 | ||
|
|
cdaa40d2a8 | ||
|
|
ca3618bc69 | ||
|
|
b2f7d2560a | ||
|
|
1ff9429655 | ||
|
|
af453e5647 | ||
|
|
32aef44388 | ||
|
|
7a74a9662b | ||
|
|
b6754f536e | ||
|
|
793cf79c89 | ||
|
|
50ac1c7bab | ||
|
|
f04d3f640e | ||
|
|
0a9396a25e | ||
|
|
038ec293b1 | ||
|
|
894ebb27f5 | ||
|
|
c9a788eedc | ||
|
|
0762f2afeb | ||
|
|
31be872f55 | ||
|
|
94c0ef3001 | ||
|
|
af1f036a70 |
@@ -18,6 +18,8 @@ steps:
|
||||
TERM: "xterm-256color"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1 # Transient Docker/BuildKit failure
|
||||
limit: 1
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
@@ -46,6 +48,8 @@ steps:
|
||||
VLLM_BRANCH: "$BUILDKITE_COMMIT"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1 # Transient Docker/BuildKit failure
|
||||
limit: 1
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
@@ -72,6 +76,8 @@ steps:
|
||||
VLLM_BRANCH: "$BUILDKITE_COMMIT"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1 # Transient Docker/BuildKit failure
|
||||
limit: 1
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
|
||||
@@ -18,6 +18,8 @@ steps:
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
- tests/kernels/mamba/test_cpu_short_conv.py
|
||||
- tests/kernels/mamba/test_causal_conv1d.py
|
||||
- tests/kernels/mamba/test_mamba_ssm.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
@@ -28,7 +30,9 @@ steps:
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# Note: SDE can't be downloaded from CI host because of AWS WAF
|
||||
# - label: CPU-Compatibility Tests
|
||||
|
||||
@@ -99,6 +99,7 @@ steps:
|
||||
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
|
||||
pytest -v -s v1/structured_output &&
|
||||
pytest -v -s v1/test_serial_utils.py &&
|
||||
pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]" &&
|
||||
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_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --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'
|
||||
- label: "XPU server test"
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# For hf script, without -t option (tensor parallel size).
|
||||
# bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
|
||||
model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
|
||||
rocm_safetensors_load_strategy: lazy
|
||||
required_gpu_arch:
|
||||
- gfx942
|
||||
- gfx950
|
||||
|
||||
@@ -72,6 +72,11 @@ def launch_lm_eval(eval_config, tp_size):
|
||||
if moe_backend is not None:
|
||||
model_args += f"moe_backend={moe_backend},"
|
||||
|
||||
if current_platform.is_rocm():
|
||||
rocm_load_strategy = eval_config.get("rocm_safetensors_load_strategy")
|
||||
if rocm_load_strategy is not None:
|
||||
model_args += f"safetensors_load_strategy={rocm_load_strategy},"
|
||||
|
||||
env_vars = eval_config.get("env_vars", None)
|
||||
with scoped_env_vars(env_vars):
|
||||
results = lm_eval.simple_evaluate(
|
||||
|
||||
+472
-399
File diff suppressed because it is too large
Load Diff
@@ -3,7 +3,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Append a build artifact line to the Buildkite annotation.
|
||||
# Usage: annotate-build-artifact.sh <label> <value>
|
||||
# Usage: annotate-build-artifact.sh <label> <value> <context>
|
||||
set -e
|
||||
echo "- **${1}**: \`${2}\`" | \
|
||||
buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
|
||||
buildkite-agent annotate --append --style 'info' \
|
||||
--context "${3:?context is required}"
|
||||
|
||||
Executable
+32
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Build the macOS arm64 CPU wheel natively on a macOS agent (the `macmini`
|
||||
# queue) into artifacts/dist/ for upload-nightly-wheels.sh.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# The Rust frontend build needs protoc.
|
||||
if ! command -v protoc >/dev/null 2>&1; then
|
||||
brew install protobuf
|
||||
fi
|
||||
|
||||
# upload-nightly-wheels.sh expects exactly one wheel.
|
||||
rm -rf artifacts/dist
|
||||
mkdir -p artifacts/dist
|
||||
|
||||
export VLLM_TARGET_DEVICE=cpu
|
||||
export VLLM_REQUIRE_RUST_FRONTEND=1
|
||||
export MACOSX_DEPLOYMENT_TARGET=11.0
|
||||
# uv's CPython is universal2; force an arm64-only build and tag so the wheel
|
||||
# isn't mislabelled universal2 and installed on Intel Macs where import fails.
|
||||
export ARCHFLAGS="-arch arm64"
|
||||
export _PYTHON_HOST_PLATFORM="macosx-11.0-arm64"
|
||||
export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}"
|
||||
|
||||
uv venv --python 3.12
|
||||
uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match
|
||||
uv build --wheel --no-build-isolation -o artifacts/dist
|
||||
|
||||
ls -l artifacts/dist/*.whl
|
||||
@@ -285,7 +285,7 @@ get_content_arg_names() {
|
||||
fi | awk 'NF && !seen[$0]++'
|
||||
}
|
||||
|
||||
compute_ci_base_content_hash() {
|
||||
compute_ci_base_content_hash_once() {
|
||||
local -a content_paths=()
|
||||
local -a content_args=()
|
||||
local dockerfile="${CI_BASE_DOCKERFILE:-}"
|
||||
@@ -301,7 +301,8 @@ compute_ci_base_content_hash() {
|
||||
if [[ -n "${dockerfile}" ]]; then
|
||||
printf 'dockerfile:%s\n' "${dockerfile}"
|
||||
printf 'resolved-build-args:\n'
|
||||
hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}"
|
||||
hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \
|
||||
|| return 1
|
||||
if [[ -n "${stages}" ]]; then
|
||||
printf 'dockerfile-stages:%s\n' "${stages}"
|
||||
if [[ -f "${dockerfile}" ]]; then
|
||||
@@ -314,6 +315,53 @@ compute_ci_base_content_hash() {
|
||||
} | sha256sum | cut -d' ' -f1
|
||||
}
|
||||
|
||||
compute_ci_base_content_hash() {
|
||||
local attempts="${CI_BASE_HASH_ATTEMPTS:-3}"
|
||||
local delay_secs="${CI_BASE_HASH_RETRY_DELAY:-5}"
|
||||
local attempt=0
|
||||
local hash=""
|
||||
local failed=0
|
||||
local -a hashes=()
|
||||
|
||||
if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ ]]; then
|
||||
echo "Invalid CI_BASE_HASH_ATTEMPTS: ${attempts}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ ! "${delay_secs}" =~ ^[0-9]+$ ]]; then
|
||||
echo "Invalid CI_BASE_HASH_RETRY_DELAY: ${delay_secs}" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
for ((attempt = 1; attempt <= attempts; attempt++)); do
|
||||
if ! hash=$(compute_ci_base_content_hash_once); then
|
||||
echo "ci_base content hash calculation ${attempt}/${attempts} failed" >&2
|
||||
failed=1
|
||||
else
|
||||
hashes+=("${hash}")
|
||||
echo "ci_base content hash calculation ${attempt}/${attempts}: ${hash}" >&2
|
||||
fi
|
||||
|
||||
if ((attempt < attempts)); then
|
||||
sleep "${delay_secs}"
|
||||
fi
|
||||
done
|
||||
|
||||
if ((failed)) || ((${#hashes[@]} != attempts)); then
|
||||
echo "Could not calculate a reliable ci_base content hash" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
for hash in "${hashes[@]:1}"; do
|
||||
if [[ "${hash}" != "${hashes[0]}" ]]; then
|
||||
echo "ci_base content hash changed between calculations" >&2
|
||||
printf ' observed: %s\n' "${hashes[@]}" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
|
||||
printf '%s\n' "${hashes[0]}"
|
||||
}
|
||||
|
||||
extract_dockerfile_arg_default() {
|
||||
local dockerfile="$1"
|
||||
local arg_name="$2"
|
||||
@@ -366,7 +414,11 @@ hash_dockerfile_arg_values() {
|
||||
printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
|
||||
if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
|
||||
digest=$(resolve_image_digest "${arg_value}")
|
||||
printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}"
|
||||
if [[ -z "${digest}" ]]; then
|
||||
echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2
|
||||
return 1
|
||||
fi
|
||||
printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}"
|
||||
fi
|
||||
done
|
||||
}
|
||||
@@ -1898,8 +1950,8 @@ confirm_remote_image_push() {
|
||||
fi
|
||||
|
||||
if [[ -z "${remote_revision}" \
|
||||
&& ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 \
|
||||
&& image_tag_is_commit_scoped ]]; then
|
||||
&& ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 ]] \
|
||||
&& image_tag_is_commit_scoped; then
|
||||
echo "Remote image exists under a commit-scoped tag; accepting push despite missing revision label."
|
||||
return 0
|
||||
fi
|
||||
@@ -2029,36 +2081,57 @@ upload_wheel_artifacts_if_present() {
|
||||
local wheel_dir="./wheel-export"
|
||||
local artifact_dir="artifacts/vllm-rocm-install"
|
||||
local archive_name="vllm-rocm-install.tar.gz"
|
||||
local metadata_dir="${wheel_dir}/.vllm-ci-artifact"
|
||||
local native_base_image=""
|
||||
local whl=""
|
||||
local whl_name=""
|
||||
local -a wheels=()
|
||||
|
||||
if ! should_upload_wheel_artifacts; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [[ ! -d "${wheel_dir}" ]] || ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
|
||||
echo "No ROCm wheel artifacts found in ${wheel_dir}"
|
||||
return 0
|
||||
if [[ -d "${wheel_dir}" ]]; then
|
||||
mapfile -t wheels < <(find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print)
|
||||
fi
|
||||
if [[ ${#wheels[@]} -ne 1 ]]; then
|
||||
echo "Expected exactly one ROCm wheel in ${wheel_dir}; found ${#wheels[@]}" >&2
|
||||
return 1
|
||||
fi
|
||||
whl="${wheels[0]}"
|
||||
whl_name=$(basename "${whl}")
|
||||
native_base_image="${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE:-}}"
|
||||
if [[ -z "${native_base_image}" ]]; then
|
||||
echo "Native ROCm artifact requires a ci_base image reference" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
echo "--- :package: Uploading ROCm vLLM install artifact"
|
||||
mkdir -p "${artifact_dir}"
|
||||
rm -rf "${artifact_dir}" "${metadata_dir}"
|
||||
mkdir -p "${artifact_dir}" "${metadata_dir}"
|
||||
|
||||
printf '%s\n' "${BUILDKITE_COMMIT:-local}" > "${metadata_dir}/commit.txt"
|
||||
printf '%s\n' "${native_base_image}" > "${metadata_dir}/native-base-image.txt"
|
||||
printf '%s\n' "${CI_BASE_IMAGE:-}" > "${metadata_dir}/ci-base-image.txt"
|
||||
printf '%s\n' "${IMAGE_TAG:-}" > "${metadata_dir}/fallback-image.txt"
|
||||
printf '%s\n' "${whl_name}" > "${metadata_dir}/wheel-filename.txt"
|
||||
|
||||
tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" .
|
||||
(
|
||||
cd "${artifact_dir}"
|
||||
sha256sum "${archive_name}" > "${archive_name}.sha256"
|
||||
)
|
||||
echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)"
|
||||
printf '%s\n' "${CI_BASE_IMAGE:-}" > "${artifact_dir}/ci-base-image.txt"
|
||||
printf '%s\n' "${IMAGE_TAG:-}" > "${artifact_dir}/fallback-image.txt"
|
||||
|
||||
for whl in "${wheel_dir}"/*.whl; do
|
||||
[[ -f "${whl}" ]] || continue
|
||||
whl_name=$(basename "${whl}")
|
||||
cp "${whl}" "${artifact_dir}/${whl_name}"
|
||||
echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
|
||||
done
|
||||
cp "${metadata_dir}"/*.txt "${artifact_dir}/"
|
||||
cp "${whl}" "${artifact_dir}/${whl_name}"
|
||||
echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)"
|
||||
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent artifact upload "${artifact_dir}/*"
|
||||
buildkite-agent artifact upload "${artifact_dir}/*" || return 1
|
||||
echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/"
|
||||
elif [[ "${BUILDKITE:-false}" == "true" ]]; then
|
||||
echo "buildkite-agent not found; cannot upload required ROCm artifacts" >&2
|
||||
return 1
|
||||
else
|
||||
echo "Not in Buildkite, skipping artifact upload"
|
||||
fi
|
||||
@@ -2090,6 +2163,11 @@ main() {
|
||||
echo "BAKE_PRINT_ONLY=1 set; skipping build"
|
||||
return 0
|
||||
fi
|
||||
if should_upload_wheel_artifacts; then
|
||||
# wheel-export is an output directory, not a BuildKit cache. Starting
|
||||
# clean prevents a failed/retried export from packaging a stale wheel.
|
||||
rm -rf ./wheel-export
|
||||
fi
|
||||
seed_dependency_caches_if_needed
|
||||
run_bake
|
||||
upload_wheel_artifacts_if_present
|
||||
|
||||
@@ -45,8 +45,10 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
|
||||
|
||||
# copy to /nightly/ only if it is on the main branch and not a PR
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
|
||||
# copy to /nightly/ only when enabled for a main branch build that is not a PR
|
||||
if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
|
||||
"$BUILDKITE_BRANCH" == "main" && \
|
||||
"$BUILDKITE_PULL_REQUEST" == "false" ]]; then
|
||||
echo "Uploading indices to overwrite /nightly/"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
|
||||
fi
|
||||
@@ -67,7 +69,7 @@ pure_version="${version%%+*}"
|
||||
echo "Pure version (without variant): $pure_version"
|
||||
|
||||
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
|
||||
if [[ "$version" != *"dev"* ]]; then
|
||||
if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
|
||||
echo "Re-generating indices for /$pure_version/"
|
||||
rm -rf "${INDICES_OUTPUT_DIR:?}"
|
||||
mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script runs tests inside the corresponding ROCm docker container.
|
||||
# It handles both single-node and multi-node test configurations.
|
||||
# This script runs ROCm tests either directly in a native CI pod or inside the
|
||||
# corresponding Docker container. Multi-node tests continue to use Docker.
|
||||
#
|
||||
# Multi-node detection: Instead of matching on fragile group names, we detect
|
||||
# multi-node jobs structurally by looking for the bracket command syntax
|
||||
@@ -34,10 +34,27 @@ set -o pipefail
|
||||
: "${CLICOLOR_FORCE:=1}"
|
||||
: "${PY_COLORS:=1}"
|
||||
: "${ROCM_DOCKER_TTY:=1}"
|
||||
: "${PYTHONFAULTHANDLER:=1}"
|
||||
: "${PYTEST_TIMEOUT:=2100}"
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
|
||||
fi
|
||||
export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations="* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations=25"
|
||||
fi
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0"
|
||||
fi
|
||||
# Dump stacks after 15 minutes, then stop an individual test after 35 minutes.
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=900"
|
||||
fi
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* &&
|
||||
" ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--timeout-method=thread"
|
||||
fi
|
||||
export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
|
||||
export PYTHONFAULTHANDLER
|
||||
|
||||
# Export Python path for commands that run directly on the host. Containerized
|
||||
# tests set this to /vllm-workspace below so spawned Python processes do not
|
||||
@@ -53,6 +70,28 @@ report_docker_usage() {
|
||||
docker system df || true
|
||||
}
|
||||
|
||||
clear_ci_orchestration_env() {
|
||||
unset -v \
|
||||
VLLM_TEST_GROUP_NAME \
|
||||
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
|
||||
VLLM_CI_ARTIFACT_STEP \
|
||||
VLLM_TEST_CACHE \
|
||||
VLLM_CI_EXECUTION_MODE \
|
||||
VLLM_CI_WORKSPACE \
|
||||
VLLM_CI_REQUIRE_WORKSPACE_MOUNT \
|
||||
VLLM_TEST_COMMANDS \
|
||||
VLLM_CI_BRANCH \
|
||||
VLLM_CI_BASE_IMAGE \
|
||||
VLLM_CI_FALLBACK_IMAGE \
|
||||
VLLM_CI_DOCKER_DISABLED \
|
||||
VLLM_CI_ARTIFACT_GLOB \
|
||||
VLLM_CI_ARTIFACT_CHECKSUM_GLOB \
|
||||
VLLM_CI_EXPECTED_GPU_COUNT \
|
||||
VLLM_CI_USE_ARTIFACTS \
|
||||
VLLM_CI_RESULTS_ROOT \
|
||||
VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
|
||||
}
|
||||
|
||||
cleanup_network() {
|
||||
local max_nodes=${NUM_NODES:-2}
|
||||
for node in $(seq 0 $((max_nodes - 1))); do
|
||||
@@ -145,7 +184,11 @@ prepare_artifact_image() {
|
||||
fi
|
||||
|
||||
cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
|
||||
tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
|
||||
tar -C "${wheel_dir}" \
|
||||
--exclude='*.whl' \
|
||||
--exclude='.vllm-ci-artifact' \
|
||||
--exclude='./.vllm-ci-artifact' \
|
||||
-cf - . \
|
||||
| tar -C "${workspace_dir}" -xf - || return 1
|
||||
cat > "${context_dir}/Dockerfile" <<'EOF'
|
||||
ARG BASE_IMAGE
|
||||
@@ -168,6 +211,257 @@ EOF
|
||||
return 0
|
||||
}
|
||||
|
||||
is_native_runtime() {
|
||||
[[ "${AMD_CI_RUNTIME:-}" == "native" || "${NATIVE_CI:-}" == "true" ]]
|
||||
}
|
||||
|
||||
validate_native_workspace() {
|
||||
local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
|
||||
local workspace_real=""
|
||||
local checkout_real=""
|
||||
local workspace_mount=""
|
||||
|
||||
mkdir -p "${workspace_dir}" || return 1
|
||||
workspace_real=$(readlink -m "${workspace_dir}") || return 1
|
||||
if [[ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ]]; then
|
||||
checkout_real=$(readlink -m "${BUILDKITE_BUILD_CHECKOUT_PATH}") || return 1
|
||||
if [[ "${checkout_real}" == "${workspace_real}" \
|
||||
|| "${checkout_real}" == "${workspace_real}/"* \
|
||||
|| "${workspace_real}" == "${checkout_real}/"* ]]; then
|
||||
echo "Refusing to replace ${workspace_real}; it overlaps the Buildkite checkout ${checkout_real}" >&2
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
if [[ "${VLLM_CI_REQUIRE_WORKSPACE_MOUNT:-1}" == "1" ]]; then
|
||||
if ! command -v findmnt >/dev/null 2>&1; then
|
||||
echo "findmnt is required to verify the native workspace mount" >&2
|
||||
return 1
|
||||
fi
|
||||
workspace_mount=$(findmnt -n -T "${workspace_real}" -o TARGET 2>/dev/null || true)
|
||||
if [[ "$(readlink -m "${workspace_mount:-/}")" != "${workspace_real}" ]]; then
|
||||
echo "Native CI requires a dedicated volume mounted at ${workspace_real}" >&2
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
||||
prepare_native_workspace() {
|
||||
if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
|
||||
echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1"
|
||||
return 1
|
||||
fi
|
||||
if ! command -v buildkite-agent >/dev/null 2>&1; then
|
||||
echo "buildkite-agent not found; cannot download ROCm wheel artifact"
|
||||
return 1
|
||||
fi
|
||||
validate_native_workspace || return 1
|
||||
|
||||
local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
|
||||
local artifact_checksum_glob="${VLLM_CI_ARTIFACT_CHECKSUM_GLOB:-${artifact_glob}.sha256}"
|
||||
local artifact_step="${VLLM_CI_ARTIFACT_STEP:-image-build-amd}"
|
||||
local archive=""
|
||||
local checksum=""
|
||||
local download_dir=""
|
||||
local metadata_dir=""
|
||||
local recorded_base=""
|
||||
local recorded_commit=""
|
||||
local recorded_wheel=""
|
||||
local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
|
||||
local wheel_dir=""
|
||||
local attempt=0
|
||||
local attempt_dir=""
|
||||
local -a archives=()
|
||||
local -a checksums=()
|
||||
local -a wheels=()
|
||||
|
||||
artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) || return 1
|
||||
wheel_dir="${artifact_work_dir}/wheels"
|
||||
mkdir -p "${wheel_dir}" || return 1
|
||||
|
||||
echo "--- Downloading ROCm wheel artifact from ${artifact_step} (native in-pod)"
|
||||
for attempt in 1 2 3; do
|
||||
attempt_dir="${artifact_work_dir}/download-${attempt}"
|
||||
rm -rf "${attempt_dir}" || return 1
|
||||
mkdir -p "${attempt_dir}" || return 1
|
||||
if buildkite-agent artifact download \
|
||||
"${artifact_glob}" "${attempt_dir}" --step "${artifact_step}" \
|
||||
&& buildkite-agent artifact download \
|
||||
"${artifact_checksum_glob}" "${attempt_dir}" --step "${artifact_step}"; then
|
||||
download_dir="${attempt_dir}"
|
||||
break
|
||||
fi
|
||||
echo "Artifact download attempt ${attempt}/3 failed"
|
||||
if [[ "${attempt}" -lt 3 ]]; then
|
||||
sleep $((attempt * 2))
|
||||
fi
|
||||
done
|
||||
if [[ -z "${download_dir}" ]]; then
|
||||
echo "Failed to download ${artifact_glob} and ${artifact_checksum_glob} from ${artifact_step}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
mapfile -t archives < <(
|
||||
find "${download_dir}" -name "vllm-rocm-install.tar.gz" -type f -print
|
||||
)
|
||||
mapfile -t checksums < <(
|
||||
find "${download_dir}" -name "vllm-rocm-install.tar.gz.sha256" -type f -print
|
||||
)
|
||||
if [[ ${#archives[@]} -ne 1 || ${#checksums[@]} -ne 1 ]]; then
|
||||
echo "Expected exactly one ROCm archive and checksum; found ${#archives[@]} archive(s) and ${#checksums[@]} checksum(s)" >&2
|
||||
return 1
|
||||
fi
|
||||
archive="${archives[0]}"
|
||||
checksum="${checksums[0]}"
|
||||
if [[ "$(dirname "${archive}")" != "$(dirname "${checksum}")" ]]; then
|
||||
echo "ROCm archive and checksum were downloaded to different directories" >&2
|
||||
return 1
|
||||
fi
|
||||
(
|
||||
cd "$(dirname "${archive}")"
|
||||
sha256sum -c "$(basename "${checksum}")"
|
||||
) || return 1
|
||||
|
||||
tar --no-same-owner -xzf "${archive}" -C "${wheel_dir}" || return 1
|
||||
mapfile -t wheels < <(
|
||||
find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print
|
||||
)
|
||||
if [[ ${#wheels[@]} -ne 1 ]]; then
|
||||
echo "ROCm artifact must contain exactly one top-level wheel; found ${#wheels[@]}" >&2
|
||||
return 1
|
||||
fi
|
||||
metadata_dir="${wheel_dir}/.vllm-ci-artifact"
|
||||
for metadata_file in commit.txt native-base-image.txt wheel-filename.txt; do
|
||||
if [[ ! -s "${metadata_dir}/${metadata_file}" ]]; then
|
||||
echo "ROCm artifact metadata is missing ${metadata_file}" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
for metadata_file in ci-base-image.txt fallback-image.txt; do
|
||||
if [[ ! -f "${metadata_dir}/${metadata_file}" ]]; then
|
||||
echo "ROCm artifact metadata is missing ${metadata_file}" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
|
||||
recorded_commit=$(tr -d '\r\n' < "${metadata_dir}/commit.txt")
|
||||
recorded_base=$(tr -d '\r\n' < "${metadata_dir}/native-base-image.txt")
|
||||
recorded_wheel=$(tr -d '\r\n' < "${metadata_dir}/wheel-filename.txt")
|
||||
if [[ -z "${BUILDKITE_COMMIT:-}" || "${recorded_commit}" != "${BUILDKITE_COMMIT}" ]]; then
|
||||
echo "ROCm artifact commit ${recorded_commit} does not match ${BUILDKITE_COMMIT:-unset}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ -z "${VLLM_CI_BASE_IMAGE:-}" || "${recorded_base}" != "${VLLM_CI_BASE_IMAGE}" ]]; then
|
||||
echo "ROCm artifact base ${recorded_base} does not match ${VLLM_CI_BASE_IMAGE:-unset}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ "${recorded_wheel}" != "$(basename "${wheels[0]}")" ]]; then
|
||||
echo "ROCm artifact wheel manifest ${recorded_wheel} does not match $(basename "${wheels[0]}")" >&2
|
||||
return 1
|
||||
fi
|
||||
for required_dir in tests .buildkite requirements; do
|
||||
if [[ ! -d "${wheel_dir}/${required_dir}" ]]; then
|
||||
echo "ROCm wheel artifact did not contain ${required_dir}/" >&2
|
||||
return 1
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- Installing ROCm wheel into pod environment"
|
||||
python3 -m pip install --no-deps --force-reinstall "${wheels[0]}" || return 1
|
||||
|
||||
echo "--- Preparing ${workspace_dir} from artifact"
|
||||
find "${workspace_dir}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + || return 1
|
||||
tar -C "${wheel_dir}" \
|
||||
--exclude='*.whl' \
|
||||
--exclude='.vllm-ci-artifact' \
|
||||
--exclude='./.vllm-ci-artifact' \
|
||||
-cf - . | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1
|
||||
if [[ ! -d "${workspace_dir}/tests" ]]; then
|
||||
echo "Failed to stage the native test workspace" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
initialize_native_environment() {
|
||||
local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
|
||||
local job_id_suffix=""
|
||||
local native_root=""
|
||||
local hf_mount=""
|
||||
|
||||
if [[ "$(id -u)" -ne 0 ]]; then
|
||||
echo "Native ROCm CI currently requires the ci_base container to run as root" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
job_id="${job_id//[^A-Za-z0-9_.-]/_}"
|
||||
job_id_suffix="${job_id##*-}"
|
||||
job_id_suffix="${job_id_suffix:0:12}"
|
||||
native_root="/tmp/vllm-native-${job_id}"
|
||||
TMPDIR="/tmp/vllm-${job_id_suffix}/tmp"
|
||||
VLLM_RPC_BASE_PATH="/tmp"
|
||||
: "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}"
|
||||
: "${TRITON_CACHE_DIR:=${native_root}/cache/triton}"
|
||||
: "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}"
|
||||
: "${XDG_CACHE_HOME:=${native_root}/cache/xdg}"
|
||||
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
|
||||
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
|
||||
: "${HF_HUB_ETAG_TIMEOUT:=60}"
|
||||
export TMPDIR VLLM_RPC_BASE_PATH
|
||||
export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME
|
||||
export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
|
||||
export PYTORCH_ROCM_ARCH=""
|
||||
|
||||
mkdir -p "${TMPDIR}" \
|
||||
"${TORCHINDUCTOR_CACHE_DIR}" \
|
||||
"${TRITON_CACHE_DIR}" \
|
||||
"${VLLM_CACHE_ROOT}" \
|
||||
"${XDG_CACHE_HOME}" \
|
||||
"${HF_HOME}" || return 1
|
||||
|
||||
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
|
||||
if ! command -v findmnt >/dev/null 2>&1; then
|
||||
echo "findmnt is required to verify the native Hugging Face cache mount" >&2
|
||||
return 1
|
||||
fi
|
||||
hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
|
||||
if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
|
||||
echo "Native CI requires a persistent volume mounted at or above ${HF_HOME}" >&2
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
||||
run_native_preflight() {
|
||||
local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
|
||||
|
||||
if [[ ! "${expected_gpus}" =~ ^[0-9]+$ ]]; then
|
||||
echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=${expected_gpus}" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
python3 -c "import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing" || return 1
|
||||
|
||||
if [[ "${expected_gpus}" == "0" ]]; then
|
||||
echo "Native CPU-only AMD job: skipping ROCm device validation"
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "--- ROCm info"
|
||||
rocminfo || return 1
|
||||
VLLM_CI_EXPECTED_GPU_COUNT="${expected_gpus}" python3 - <<'PY'
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
expected = int(os.environ["VLLM_CI_EXPECTED_GPU_COUNT"])
|
||||
assert torch.version.hip, "PyTorch is not a ROCm build"
|
||||
assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
|
||||
actual = torch.cuda.device_count()
|
||||
assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
|
||||
PY
|
||||
}
|
||||
|
||||
is_multi_node() {
|
||||
local cmds="$1"
|
||||
# Primary signal: NUM_NODES environment variable set by the pipeline
|
||||
@@ -350,7 +644,58 @@ re_quote_pytest_markers() {
|
||||
# Main
|
||||
###############################################################################
|
||||
|
||||
# --- GPU initialization ---
|
||||
if is_native_runtime; then
|
||||
echo "--- Native in-pod ROCm CI (AMD_CI_RUNTIME=${AMD_CI_RUNTIME:-unset}, NATIVE_CI=${NATIVE_CI:-unset})"
|
||||
artifact_work_dir=""
|
||||
|
||||
cleanup_native_workspace() {
|
||||
if [[ -n "${artifact_work_dir}" ]]; then
|
||||
rm -rf "${artifact_work_dir}"
|
||||
fi
|
||||
}
|
||||
trap cleanup_native_workspace EXIT
|
||||
|
||||
if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
|
||||
commands="${VLLM_TEST_COMMANDS}"
|
||||
commands_source="env"
|
||||
else
|
||||
commands="$*"
|
||||
commands_source="argv"
|
||||
if [[ -z "$commands" ]]; then
|
||||
echo "Error: No test commands provided for native CI." >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ "$commands_source" == "argv" ]]; then
|
||||
commands=$(re_quote_pytest_markers "$commands")
|
||||
fi
|
||||
|
||||
if is_multi_node "$commands"; then
|
||||
echo "Native CI does not support multi-node jobs yet."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! initialize_native_environment; then
|
||||
echo "Failed to initialize the native test environment"
|
||||
exit 1
|
||||
fi
|
||||
if ! prepare_native_workspace; then
|
||||
echo "Failed to prepare native test workspace"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export PYTHONPATH="${VLLM_CI_WORKSPACE:-/vllm-workspace}"
|
||||
|
||||
echo "Native test commands: $commands"
|
||||
run_native_preflight || exit 1
|
||||
# Keep AMD CI orchestration variables out of vLLM's runtime environment.
|
||||
clear_ci_orchestration_env
|
||||
/bin/bash -o pipefail -c "${commands}"
|
||||
handle_pytest_exit "$?"
|
||||
fi
|
||||
|
||||
# --- GPU initialization for legacy Docker execution ---
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
@@ -452,25 +797,27 @@ fi
|
||||
|
||||
echo "Final commands: $commands"
|
||||
|
||||
# The ROCm test image often ships /vllm-workspace without .git (artifact tarball unpack).
|
||||
# tests/standalone_tests/python_only_compile.sh uses merge-base(HEAD, origin/main) for
|
||||
# wheels.vllm.ai; compute on the agent (full git checkout) and pass into the container.
|
||||
vllm_standalone_merge_base=""
|
||||
checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
|
||||
if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
|
||||
checkout="."
|
||||
standalone_merge_base_env=()
|
||||
if [[ "$commands" == *python_only_compile.sh* ]]; then
|
||||
# The ROCm test image often ships /vllm-workspace without .git. Resolve the
|
||||
# wheels.vllm.ai commit from the agent checkout for this test only.
|
||||
vllm_standalone_merge_base=""
|
||||
checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
|
||||
if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
|
||||
checkout="."
|
||||
fi
|
||||
# Pass safe.directory per-command because Buildkite uses mixed user IDs.
|
||||
if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
|
||||
vllm_standalone_merge_base="$(
|
||||
git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
|
||||
)"
|
||||
fi
|
||||
if [[ -z "${vllm_standalone_merge_base}" ]]; then
|
||||
vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
|
||||
fi
|
||||
echo "INFO: passing CI_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
|
||||
standalone_merge_base_env=(-e "CI_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}")
|
||||
fi
|
||||
# Pass safe.directory per-command (-c) because buildkite runs will always fail
|
||||
# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
|
||||
if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
|
||||
vllm_standalone_merge_base="$(
|
||||
git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
|
||||
)"
|
||||
fi
|
||||
if [[ -z "${vllm_standalone_merge_base}" ]]; then
|
||||
vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}"
|
||||
fi
|
||||
echo "INFO: passing VLLM_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}"
|
||||
|
||||
MYPYTHONPATH="/vllm-workspace"
|
||||
|
||||
@@ -501,6 +848,7 @@ else
|
||||
fi
|
||||
|
||||
# --- Route: multi-node vs single-node ---
|
||||
clear_ci_orchestration_env
|
||||
if is_multi_node "$commands"; then
|
||||
echo "--- Multi-node job detected"
|
||||
export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
|
||||
@@ -589,7 +937,9 @@ else
|
||||
-e FORCE_COLOR \
|
||||
-e CLICOLOR_FORCE \
|
||||
-e PY_COLORS \
|
||||
-e PYTHONFAULTHANDLER \
|
||||
-e PYTEST_ADDOPTS \
|
||||
-e PYTEST_TIMEOUT \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
@@ -599,7 +949,7 @@ else
|
||||
-e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
|
||||
-e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
|
||||
-e "PYTORCH_ROCM_ARCH=" \
|
||||
-e "VLLM_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}" \
|
||||
"${standalone_merge_base_env[@]}" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
|
||||
|
||||
@@ -40,7 +40,9 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
@@ -97,3 +99,4 @@ function cpu_tests() {
|
||||
# All of CPU tests are expected to be finished less than 40 mins.
|
||||
export -f cpu_tests
|
||||
timeout 2h bash -c cpu_tests
|
||||
|
||||
|
||||
@@ -35,6 +35,7 @@ case "${test_suite}" in
|
||||
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
|
||||
pytest -v -s v1/structured_output
|
||||
pytest -v -s v1/test_serial_utils.py
|
||||
pytest -v -s v1/e2e/general/test_correctness_sliding_window.py --deselect="tests/v1/e2e/general/test_correctness_sliding_window.py::test_sliding_window_retrieval[True-1-5-google/gemma-3-1b-it]"
|
||||
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_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --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
|
||||
;;
|
||||
|
||||
@@ -13,6 +13,18 @@ metadata_get() {
|
||||
fi
|
||||
}
|
||||
|
||||
use_ci_base_if_present() {
|
||||
local ci_base_image=""
|
||||
|
||||
ci_base_image="$(metadata_get rocm-ci-base-image)"
|
||||
if [[ -z "${ci_base_image}" ]]; then
|
||||
return 1
|
||||
fi
|
||||
|
||||
export CI_BASE_IMAGE="${ci_base_image}"
|
||||
echo "Using ROCm ci_base image selected by the preceding build step: ${CI_BASE_IMAGE}"
|
||||
}
|
||||
|
||||
use_refreshed_base_if_present() {
|
||||
local base_refreshed=""
|
||||
|
||||
@@ -22,15 +34,12 @@ use_refreshed_base_if_present() {
|
||||
fi
|
||||
|
||||
export BASE_IMAGE
|
||||
export CI_BASE_IMAGE
|
||||
export IMAGE_TAG_LATEST
|
||||
|
||||
BASE_IMAGE="$(metadata_get rocm-base-image)"
|
||||
CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
|
||||
IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"
|
||||
|
||||
echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
|
||||
echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
|
||||
if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
|
||||
echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
|
||||
fi
|
||||
@@ -41,6 +50,8 @@ use_refreshed_base_if_present() {
|
||||
main() {
|
||||
local base_refreshed=0
|
||||
|
||||
use_ci_base_if_present || true
|
||||
|
||||
if use_refreshed_base_if_present; then
|
||||
base_refreshed=1
|
||||
fi
|
||||
|
||||
@@ -6,8 +6,14 @@ set -ex
|
||||
# manylinux platform tag with auditwheel.
|
||||
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
|
||||
|
||||
# shellcheck source=lib/manylinux.sh
|
||||
source .buildkite/scripts/lib/manylinux.sh
|
||||
# auditwheel is Linux-only; macOS wheels already carry a valid tag, so skip the
|
||||
# manylinux retag for them.
|
||||
WHEEL_PLATFORM="${VLLM_WHEEL_PLATFORM:-linux}"
|
||||
|
||||
if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
|
||||
# shellcheck source=lib/manylinux.sh
|
||||
source .buildkite/scripts/lib/manylinux.sh
|
||||
fi
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
@@ -27,8 +33,10 @@ wheel="${wheel_files[0]}"
|
||||
|
||||
# ========= detect manylinux tag and rename ==========
|
||||
|
||||
wheel="$(apply_manylinux_tag "$wheel")"
|
||||
echo "Renamed wheel to: $wheel"
|
||||
if [[ "$WHEEL_PLATFORM" == "linux" ]]; then
|
||||
wheel="$(apply_manylinux_tag "$wheel")"
|
||||
echo "Renamed wheel to: $wheel"
|
||||
fi
|
||||
|
||||
# Extract the version from the wheel
|
||||
version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
|
||||
|
||||
@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
|
||||
|
||||
# Update rocm/nightly/ if on main branch and not a PR
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
|
||||
# Only scheduled nightly builds should update the moving nightly index.
|
||||
if [[ "${NIGHTLY:-0}" == "1" ]]; then
|
||||
echo "Updating rocm/nightly/ index..."
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
|
||||
fi
|
||||
@@ -147,7 +147,7 @@ echo ""
|
||||
echo "Install command (by commit):"
|
||||
echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
|
||||
echo ""
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
|
||||
if [[ "${NIGHTLY:-0}" == "1" ]]; then
|
||||
echo "Install command (nightly):"
|
||||
echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
|
||||
fi
|
||||
|
||||
+160
-28
File diff suppressed because it is too large
Load Diff
@@ -13,6 +13,7 @@ steps:
|
||||
- pytest -v -s benchmarks/
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -18,6 +18,7 @@ steps:
|
||||
- pytest -v -s cuda/test_platform_no_cuda_init.py
|
||||
|
||||
- label: Cudagraph
|
||||
device: h200_35gb
|
||||
key: cudagraph
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
@@ -28,4 +29,4 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
|
||||
@@ -15,6 +15,7 @@ steps:
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 85
|
||||
depends_on:
|
||||
@@ -65,6 +66,7 @@ steps:
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
@@ -90,6 +92,7 @@ steps:
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 85
|
||||
depends_on:
|
||||
@@ -115,6 +118,7 @@ steps:
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_4
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
@@ -171,6 +175,7 @@ steps:
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
timeout_in_minutes: 70
|
||||
depends_on:
|
||||
@@ -182,7 +187,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
|
||||
- KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
|
||||
@@ -198,3 +203,25 @@ steps:
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
|
||||
# P TP 4 - D DPEP 4 test case for DSv4-Flash
|
||||
- label: DSv4-Flash Disaggregated DP EP
|
||||
key: dsv4-flash-disaggregated
|
||||
timeout_in_minutes: 60
|
||||
device: h200
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 8
|
||||
env:
|
||||
ENABLE_HMA_FLAG: "1"
|
||||
DP_EP: "1"
|
||||
GPU_MEMORY_UTILIZATION: "0.85"
|
||||
PREFILLER_TP_SIZE: "4"
|
||||
DECODER_TP_SIZE: "4"
|
||||
PREFILL_BLOCK_SIZE: "256"
|
||||
DECODE_BLOCK_SIZE: "256"
|
||||
MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
|
||||
VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
|
||||
|
||||
@@ -39,6 +39,7 @@ steps:
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -114,6 +114,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Entrypoints Unit Tests
|
||||
device: h200_35gb
|
||||
key: entrypoints-unit-tests
|
||||
timeout_in_minutes: 25
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -15,6 +16,7 @@ steps:
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-llm
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -55,6 +57,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -73,6 +76,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -92,6 +96,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-generate
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -114,6 +119,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -148,6 +154,7 @@ steps:
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -16,6 +16,7 @@ steps:
|
||||
- pytest -v -s distributed/test_eplb_utils.py
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -15,6 +15,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/ir
|
||||
|
||||
- label: Kernels Core Operation Test
|
||||
device: h200_35gb
|
||||
key: kernels-core-operation-test
|
||||
timeout_in_minutes: 120
|
||||
source_file_dependencies:
|
||||
@@ -163,6 +164,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Kernels Mamba Test
|
||||
device: h200_35gb
|
||||
key: kernels-mamba-test
|
||||
timeout_in_minutes: 40
|
||||
source_file_dependencies:
|
||||
@@ -176,9 +178,9 @@ steps:
|
||||
timeout_in_minutes: 25
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/fla/ops/kda.py
|
||||
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
|
||||
- vllm/model_executor/layers/fla/ops/l2norm.py
|
||||
- vllm/third_party/flash_linear_attention/ops/kda.py
|
||||
- vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
|
||||
- vllm/third_party/flash_linear_attention/ops/l2norm.py
|
||||
- tests/kernels/test_kda.py
|
||||
commands:
|
||||
- pytest -v -s kernels/test_kda.py
|
||||
@@ -232,6 +234,15 @@ steps:
|
||||
- vllm/v1/attention/backends/mla/flashinfer_mla.py
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/platforms/cuda.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
|
||||
- vllm/cute_utils/
|
||||
- vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
|
||||
- vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
|
||||
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
|
||||
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
|
||||
- tests/kernels/test_ll_bf16_gemm.py
|
||||
- tests/kernels/test_top_k_per_row.py
|
||||
commands:
|
||||
- nvidia-smi
|
||||
@@ -260,6 +271,9 @@ steps:
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
|
||||
- pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
|
||||
- pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
|
||||
- pytest -v -s tests/kernels/test_ll_bf16_gemm.py
|
||||
# e2e
|
||||
- pytest -v -s tests/models/quantization/test_nvfp4.py
|
||||
|
||||
|
||||
@@ -78,6 +78,28 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
|
||||
|
||||
- label: LM Eval PCP (4xB200)
|
||||
key: lm-eval-pcp-4xb200
|
||||
timeout_in_minutes: 360
|
||||
device: b200-k8s
|
||||
num_devices: 4
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
|
||||
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
|
||||
- tests/evals/gsm8k/configs/models-pcp.txt
|
||||
- vllm/model_executor/layers/quantization
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/parallel_state.py
|
||||
- vllm/model_executor/layers/attention/mla_attention.py
|
||||
- vllm/model_executor/layers/attention/pcp.py
|
||||
- vllm/v1/worker/gpu/model_runner.py
|
||||
- vllm/v1/worker/gpu/pcp_manager.py
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
|
||||
|
||||
- label: LM Eval Large Models EP (2xB200)
|
||||
key: lm-eval-large-models-ep-2xb200
|
||||
timeout_in_minutes: 60
|
||||
@@ -103,7 +125,7 @@ steps:
|
||||
- vllm/transformers_utils/configs/qwen3_5_moe.py
|
||||
- vllm/model_executor/models/qwen3_next.py
|
||||
- vllm/model_executor/models/qwen3_next_mtp.py
|
||||
- vllm/model_executor/layers/fla/ops/
|
||||
- vllm/third_party/flash_linear_attention/ops/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
|
||||
|
||||
@@ -117,6 +139,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_8
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
|
||||
@@ -23,6 +23,7 @@ steps:
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 75
|
||||
depends_on:
|
||||
@@ -63,8 +64,9 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Core + KV + Metrics
|
||||
device: h200_35gb
|
||||
key: v1-core-kv-metrics
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 80
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
@@ -25,6 +26,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -15,11 +15,13 @@ steps:
|
||||
- pytest -v -s models/language -m 'core_model and (not slow_test)'
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Language Models Tests (Extra Standard) %N
|
||||
device: h200_35gb
|
||||
key: language-models-tests-extra-standard
|
||||
timeout_in_minutes: 40
|
||||
source_file_dependencies:
|
||||
@@ -35,6 +37,7 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -49,8 +52,8 @@ steps:
|
||||
- tests/models/language/pooling/test_classification.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Language Models Tests (Hybrid) %N
|
||||
device: h200_35gb
|
||||
key: language-models-tests-hybrid
|
||||
timeout_in_minutes: 65
|
||||
source_file_dependencies:
|
||||
@@ -61,8 +64,8 @@ steps:
|
||||
# Note: also needed to run plamo2 model in vLLM
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
# Shard hybrid language model tests
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
# Shard the hybrid language model tests that are numerically stable on Hopper.
|
||||
- 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
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
@@ -75,6 +78,20 @@ steps:
|
||||
- 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
|
||||
|
||||
# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
|
||||
# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
|
||||
# correctness test on L4 until its H200 output matches the Transformers reference.
|
||||
- label: Language Models Tests (Granite L4 Compatibility)
|
||||
key: language-models-tests-granite-l4-compatibility
|
||||
timeout_in_minutes: 65
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
|
||||
|
||||
- label: Language Models Test (Extended Generation) # 80min
|
||||
device: h200_35gb
|
||||
key: language-models-test-extended-generation
|
||||
@@ -101,10 +118,10 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s models/language/generation_ppl_test
|
||||
|
||||
- label: Language Models Test (Extended Pooling) # 36min
|
||||
- label: Language Models Test (Extended Pooling)
|
||||
device: h200_35gb
|
||||
key: language-models-test-extended-pooling
|
||||
timeout_in_minutes: 70
|
||||
timeout_in_minutes: 120
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -114,7 +131,7 @@ steps:
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 100
|
||||
timeout_in_minutes: 120
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -107,6 +107,7 @@ steps:
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -118,6 +119,7 @@ steps:
|
||||
- vllm/model_executor/model_loader/
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 1)
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-extended-generation-1
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -109,14 +109,16 @@ steps:
|
||||
- pytest -s -v compile/passes --ignore compile/passes/distributed
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
device: h200_35gb
|
||||
key: pytorch-fullgraph-smoke-test
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -148,7 +150,42 @@ steps:
|
||||
# as it is a heavy test that is covered in other steps.
|
||||
# Use `find` to launch multiple instances of pytest so that
|
||||
# they do not suffer from https://github.com/vllm-project/vllm/issues/28965
|
||||
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
|
||||
# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
|
||||
# those SM90 cases skip while the architecture-compatible cases still run.
|
||||
- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
|
||||
key: pytorch-fullgraph-cudagraph-l4-compatibility
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/_custom_ops.py
|
||||
- vllm/compilation/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- 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
|
||||
commands:
|
||||
- pytest -s -v compile/fullgraph/test_full_cudagraph.py
|
||||
|
||||
- label: PyTorch Fullgraph
|
||||
key: pytorch-fullgraph
|
||||
@@ -197,6 +234,7 @@ steps:
|
||||
- bash standalone_tests/pytorch_nightly_dependency.sh
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -3,8 +3,11 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Quantization
|
||||
device: h200_35gb
|
||||
key: quantization
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 75
|
||||
env:
|
||||
VLLM_USE_V2_MODEL_RUNNER: "0"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -19,9 +22,13 @@ steps:
|
||||
# TODO(jerryzh168): resolve the above comment
|
||||
- uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
|
||||
- uv pip install --system conch-triton-kernels
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
|
||||
# The SM90-only checkpoint currently contains a removed weight_chan_scale
|
||||
# parameter. It was not exercised by the previous L4 job.
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 8
|
||||
|
||||
- label: Quantized Fusions
|
||||
device: h200_35gb
|
||||
key: quantized-fusions
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
@@ -52,10 +59,14 @@ steps:
|
||||
- pytest -s -v tests/quantization/test_blackwell_moe.py
|
||||
|
||||
- label: Quantized Models Test
|
||||
device: h200_35gb
|
||||
key: quantized-models-test
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 65
|
||||
env:
|
||||
VLLM_USE_V2_MODEL_RUNNER: "0"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/models/quantization
|
||||
commands:
|
||||
- pytest -v -s models/quantization
|
||||
- pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 3
|
||||
|
||||
@@ -81,6 +81,7 @@ steps:
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: Rust Frontend Tool Use
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 25
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -15,6 +15,7 @@ steps:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
|
||||
mirror:
|
||||
amd:
|
||||
dind: false
|
||||
device: mi300_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
+1
-1
@@ -172,7 +172,7 @@ mkdocs.yaml @hmellor
|
||||
# Kernels
|
||||
/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
|
||||
/vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
|
||||
/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy
|
||||
/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy
|
||||
|
||||
# ROCm related: specify owner with write access to notify AMD folks for careful code review
|
||||
/vllm/**/*rocm* @tjtanaa @dllehr-amd
|
||||
|
||||
@@ -323,7 +323,7 @@ jobs:
|
||||
// {users} will be replaced with @mentions
|
||||
const ccConfig = {
|
||||
rocm: {
|
||||
users: ['hongxiayang', 'tjtanaa', 'vllmellm'],
|
||||
users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'],
|
||||
message: 'CC {users} for ROCm-related issue',
|
||||
},
|
||||
mistral: {
|
||||
|
||||
@@ -18,6 +18,9 @@ vllm/third_party/deep_gemm/
|
||||
# fmha_sm100 vendored package built from source
|
||||
vllm/third_party/fmha_sm100/
|
||||
|
||||
# tml-fa4 vendored package built from source
|
||||
vllm/third_party/tml_fa4/
|
||||
|
||||
# triton jit
|
||||
.triton
|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ repos:
|
||||
- id: markdownlint-cli2
|
||||
language_version: lts
|
||||
args: [--fix]
|
||||
exclude: ^CLAUDE\.md$
|
||||
exclude: (^|/)CLAUDE\.md$
|
||||
- repo: https://github.com/rhysd/actionlint
|
||||
rev: v1.7.7
|
||||
hooks:
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
collect_env.py
|
||||
vllm/model_executor/layers/fla/ops/*.py
|
||||
@@ -1408,6 +1408,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
include(cmake/external_projects/fmha_sm100.cmake)
|
||||
include(cmake/external_projects/flashmla.cmake)
|
||||
include(cmake/external_projects/qutlass.cmake)
|
||||
include(cmake/external_projects/tml_fa4.cmake)
|
||||
|
||||
# vllm-flash-attn should be last as it overwrites some CMake functions
|
||||
include(cmake/external_projects/vllm_flash_attn.cmake)
|
||||
|
||||
@@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size):
|
||||
# Output workspace
|
||||
dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
|
||||
workspace_starts_t = torch.tensor(
|
||||
workspace_starts, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
return cache, dst, block_table, seq_lens_t, workspace_starts_t
|
||||
return cache, dst, block_table, workspace_starts_t
|
||||
|
||||
|
||||
def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
@@ -94,7 +93,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
)
|
||||
)
|
||||
def bench_fn(total_tokens, provider, num_reqs):
|
||||
cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
|
||||
cache, dst, block_table, ws_starts = make_inputs(
|
||||
total_tokens, num_reqs, BLOCK_SIZE
|
||||
)
|
||||
|
||||
@@ -102,7 +101,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
|
||||
cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
|
||||
cache, dst, block_table, ws_starts, num_reqs
|
||||
),
|
||||
quantiles=quantiles,
|
||||
rep=500,
|
||||
|
||||
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
|
||||
"csrc/cpu/layernorm.cpp"
|
||||
"csrc/cpu/mla_decode.cpp"
|
||||
"csrc/cpu/pos_encoding.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp")
|
||||
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
"csrc/cpu/activation.cpp"
|
||||
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a
|
||||
GIT_TAG a8f794d1251cbfd88a5011445dd5582289c727e4
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
@@ -35,7 +35,7 @@ set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
|
||||
file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
|
||||
file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
|
||||
FLASHMLA_INTERFACE_CONTENT)
|
||||
string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
|
||||
string(REPLACE "flash_mla_cuda = torch.ops._flashmla_C"
|
||||
"import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
|
||||
FLASHMLA_INTERFACE_CONTENT
|
||||
"${FLASHMLA_INTERFACE_CONTENT}")
|
||||
@@ -72,7 +72,7 @@ if(FLASH_MLA_ARCHS)
|
||||
list(APPEND VLLM_FLASHMLA_GPU_FLAGS "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math")
|
||||
|
||||
set(FlashMLA_SOURCES
|
||||
${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
|
||||
${flashmla_SOURCE_DIR}/csrc/api/api.cpp
|
||||
|
||||
# Misc kernels for decoding
|
||||
${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
|
||||
@@ -128,6 +128,7 @@ if(FLASH_MLA_ARCHS)
|
||||
|
||||
set(FlashMLA_Extension_INCLUDES
|
||||
${flashmla_SOURCE_DIR}/csrc
|
||||
${flashmla_SOURCE_DIR}/csrc/kerutils/include
|
||||
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/include
|
||||
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
|
||||
@@ -152,15 +153,18 @@ if(FLASH_MLA_ARCHS)
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Keep Stable ABI for the module, but *not* for CUDA/C++ files.
|
||||
# This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
|
||||
# Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
|
||||
# Enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
|
||||
target_compile_options(_flashmla_C PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
|
||||
|
||||
# _flashmla_C is now ABI-stable torch 2.11+
|
||||
target_compile_definitions(_flashmla_C PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
target_compile_definitions(_flashmla_C PRIVATE USE_CUDA)
|
||||
endif()
|
||||
|
||||
define_extension_target(
|
||||
_flashmla_extension_C
|
||||
DESTINATION vllm
|
||||
@@ -172,11 +176,12 @@ if(FLASH_MLA_ARCHS)
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Keep Stable ABI for the module, but *not* for CUDA/C++ files.
|
||||
# This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
|
||||
target_compile_options(_flashmla_extension_C PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
|
||||
# _flashmla_extension_C is now ABI-stable w/ torch 2.11+
|
||||
target_compile_definitions(_flashmla_extension_C PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
target_compile_definitions(_flashmla_extension_C PRIVATE USE_CUDA)
|
||||
endif()
|
||||
else()
|
||||
message(STATUS "FlashMLA will not compile: unsupported CUDA architecture ${CUDA_ARCHS}")
|
||||
# Create empty targets for setup.py on unsupported systems
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
include(FetchContent)
|
||||
|
||||
if(DEFINED ENV{TML_FA4_SRC_DIR})
|
||||
set(TML_FA4_SRC_DIR $ENV{TML_FA4_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(TML_FA4_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
tml_fa4
|
||||
SOURCE_DIR ${TML_FA4_SRC_DIR}
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND "")
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
tml_fa4
|
||||
GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
|
||||
GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND "")
|
||||
endif()
|
||||
|
||||
FetchContent_GetProperties(tml_fa4)
|
||||
if(NOT tml_fa4_POPULATED)
|
||||
FetchContent_Populate(tml_fa4)
|
||||
endif()
|
||||
message(STATUS "tml-fa4 is available at ${tml_fa4_SOURCE_DIR}")
|
||||
|
||||
add_custom_target(tml_fa4)
|
||||
|
||||
# Install into a private namespace so this implementation cannot shadow the
|
||||
# flash_attn package used by vLLM's standard attention backends.
|
||||
install(CODE "
|
||||
file(GLOB_RECURSE TML_FA4_PY_FILES
|
||||
\"${tml_fa4_SOURCE_DIR}/flash_attn/cute/*.py\")
|
||||
foreach(SRC_FILE \${TML_FA4_PY_FILES})
|
||||
file(RELATIVE_PATH REL_PATH
|
||||
\"${tml_fa4_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
|
||||
set(DST_FILE
|
||||
\"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tml_fa4/\${REL_PATH}\")
|
||||
get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
|
||||
file(MAKE_DIRECTORY \${DST_DIR})
|
||||
file(READ \${SRC_FILE} FILE_CONTENTS)
|
||||
string(REPLACE
|
||||
\"flash_attn.cute\"
|
||||
\"vllm.third_party.tml_fa4\"
|
||||
FILE_CONTENTS \"\${FILE_CONTENTS}\")
|
||||
file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
|
||||
endforeach()
|
||||
" COMPONENT tml_fa4)
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42
|
||||
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
+1
-2
@@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
|
||||
torch::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::Tensor const& seq_lens, // [BATCH]
|
||||
torch::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size);
|
||||
int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
|
||||
|
||||
// Indexer K quantization and cache function
|
||||
void indexer_k_quant_and_cache(
|
||||
|
||||
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
reg.val[1] = fp16_to_fp32_bits(raw_lo);
|
||||
}
|
||||
float reduce_sum() const {
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
float result = 0;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, &ar](int i) { result += ar.values[i]; });
|
||||
|
||||
return result;
|
||||
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
|
||||
// Step 1: pairwise sum of the two 4-wide halves
|
||||
__vector float s = vec_add(reg.val[0], reg.val[1]);
|
||||
// Step 2: rotate by 8 bytes (2 floats) and add
|
||||
s = vec_add(s, vec_sld(s, s, 8));
|
||||
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
|
||||
s = vec_add(s, vec_sld(s, s, 4));
|
||||
return vec_extract(s, 0);
|
||||
}
|
||||
FP32Vec8 exp() const {
|
||||
f32x4x2_t out;
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
|
||||
// mamba_kernels.hpp.
|
||||
|
||||
#include "cpu/mamba_kernels.hpp"
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <torch/library.h>
|
||||
#include <c10/util/Optional.h>
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update
|
||||
// ---------------------------------------------------------------------------
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
|
||||
bool do_silu = false;
|
||||
if (activation.has_value()) {
|
||||
const std::string& act = activation.value();
|
||||
do_silu = (act == "silu" || act == "swish");
|
||||
}
|
||||
|
||||
at::ScalarType dtype = x.scalar_type();
|
||||
|
||||
// Input x: contiguous in native dtype.
|
||||
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
|
||||
|
||||
// conv_state: NEVER copy the full paged tensor just for layout reasons.
|
||||
// If the dtype matches we work directly on conv_state (contiguous or not)
|
||||
// by extracting strides and passing them to the kernel.
|
||||
// Only a dtype-conversion copy is made when types differ (rare for BF16).
|
||||
bool state_type_ok = (conv_state.scalar_type() == dtype);
|
||||
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
|
||||
// state_c and conv_state may be non-contiguous — that is intentional.
|
||||
|
||||
// Weight: coerce to same dtype if needed (should match in practice)
|
||||
at::Tensor w_c =
|
||||
(weight.scalar_type() != dtype)
|
||||
? weight.to(dtype).contiguous()
|
||||
: (weight.is_contiguous() ? weight : weight.contiguous());
|
||||
|
||||
// Bias stays float32 (small scalar, used only for fp32 accumulation)
|
||||
at::Tensor bias_f32;
|
||||
if (bias.has_value() && bias.value().defined())
|
||||
bias_f32 = bias.value().to(at::kFloat).contiguous();
|
||||
|
||||
int64_t batch = x_c.size(0);
|
||||
int64_t dim = x_c.size(1);
|
||||
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
|
||||
int64_t width = w_c.size(1);
|
||||
int64_t state_len = state_c.size(2);
|
||||
|
||||
// Extract strides — works for contiguous AND non-contiguous (transposed)
|
||||
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
|
||||
// contiguous) stride(1): between conv channels (dim stride) stride(2):
|
||||
// between state elements (=1 when contiguous, =dim when transposed)
|
||||
int64_t stride_s_slot = state_c.stride(0);
|
||||
int64_t stride_s_dim = state_c.stride(1);
|
||||
int64_t stride_s_state = state_c.stride(2);
|
||||
|
||||
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
|
||||
|
||||
const int32_t* cache_idx_ptr = nullptr;
|
||||
at::Tensor cache_idx_int;
|
||||
if (conv_state_indices.has_value()) {
|
||||
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
|
||||
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
|
||||
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
|
||||
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
|
||||
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
|
||||
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cache_idx_ptr,
|
||||
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
|
||||
do_silu);
|
||||
});
|
||||
|
||||
// Write back only when a type-conversion copy was made.
|
||||
// Layout-only non-contiguity is handled via strides above — no copy needed.
|
||||
if (!state_type_ok) conv_state.copy_(state_c);
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
// ---------------------------------------------------------------------------
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, // (nstates, nheads, dim, dstate)
|
||||
const at::Tensor& x, // (N, nheads, dim)
|
||||
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C, const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens) {
|
||||
at::ScalarType state_type = state.scalar_type();
|
||||
at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
// x, B, C must be contiguous and match input_type
|
||||
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_input(x);
|
||||
at::Tensor B_in = ensure_input(B);
|
||||
at::Tensor C_in = ensure_input(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
|
||||
|
||||
// A, D, dt_bias are float32 model parameters that arrive here as expanded
|
||||
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
|
||||
// We need just the scalar value per head as a (nheads,) 1-D array so that
|
||||
// A_ptr[h] in the kernel correctly reads head h's value.
|
||||
//
|
||||
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
|
||||
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
|
||||
// → .select(2,0) → (nheads, head_dim) strides (1,0)
|
||||
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
|
||||
// No allocation, no type conversion (A is already float32).
|
||||
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
|
||||
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
|
||||
at::Tensor D_f32, dt_bias_f32;
|
||||
if (D.has_value() && D.value().defined())
|
||||
D_f32 = to_per_head_1d_f32(D.value());
|
||||
if (dt_bias.has_value() && dt_bias.value().defined())
|
||||
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
|
||||
|
||||
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
|
||||
// the type conversion so we convert head_dim x fewer elements.
|
||||
at::Tensor dt_f32;
|
||||
{
|
||||
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
|
||||
// take a zero-copy view of index 0 along that dim first.
|
||||
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
|
||||
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
|
||||
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
|
||||
}
|
||||
|
||||
int64_t nheads = state.size(1);
|
||||
int64_t dim = state.size(2);
|
||||
int64_t dstate = state.size(3);
|
||||
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
|
||||
? cu_seqlens.value().size(0) - 1
|
||||
: x_in.size(0);
|
||||
int64_t ngroups = B_in.size(1);
|
||||
|
||||
// Strides
|
||||
int64_t stride_state_n = state.stride(0);
|
||||
int64_t stride_state_h = state.stride(1);
|
||||
int64_t stride_state_d = state.stride(2);
|
||||
int64_t stride_x_n = x_in.stride(0);
|
||||
int64_t stride_x_h = x_in.stride(1);
|
||||
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
|
||||
int64_t stride_BC_n = B_in.stride(0);
|
||||
int64_t stride_BC_g = B_in.stride(1);
|
||||
int64_t stride_out_n = out.stride(0);
|
||||
int64_t stride_out_h = out.stride(1);
|
||||
|
||||
// Optional index pointers
|
||||
auto get_int32_ptr =
|
||||
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
|
||||
return (opt.has_value() && opt.value().defined())
|
||||
? opt.value().data_ptr<int32_t>()
|
||||
: nullptr;
|
||||
};
|
||||
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
|
||||
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
|
||||
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
|
||||
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
|
||||
|
||||
// Dispatch on (state_t, input_t, out_t): write directly into `out`
|
||||
// without any intermediate float32 buffer.
|
||||
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
|
||||
using state_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
|
||||
using input_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
|
||||
using out_t = scalar_t;
|
||||
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
|
||||
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
|
||||
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
|
||||
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
|
||||
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
|
||||
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
|
||||
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
|
||||
nheads, ngroups, dim, dstate, dt_softplus);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd_cpu
|
||||
// ---------------------------------------------------------------------------
|
||||
void mamba_chunk_scan_fwd_cpu_impl(
|
||||
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
|
||||
at::Tensor&
|
||||
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
|
||||
const at::Tensor& x, // [seqlen, nheads, headdim]
|
||||
const at::Tensor&
|
||||
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
|
||||
const at::Tensor& A, // [nheads] float32
|
||||
const at::Tensor& B, // [seqlen, ngroups, dstate]
|
||||
const at::Tensor& C, // [seqlen, ngroups, dstate]
|
||||
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
|
||||
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
|
||||
const at::Tensor& cu_seqlens // [batch+1] int32
|
||||
) {
|
||||
const at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_contig(x);
|
||||
at::Tensor B_in = ensure_contig(B);
|
||||
at::Tensor C_in = ensure_contig(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
|
||||
|
||||
// A and D are float32 model parameters, potentially broadcast-expanded.
|
||||
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
|
||||
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor A_f32 = to_per_head_f32(A);
|
||||
at::Tensor D_f32;
|
||||
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
|
||||
|
||||
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
|
||||
// Python.
|
||||
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
|
||||
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
|
||||
|
||||
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
|
||||
|
||||
const int64_t batch = final_states.size(0);
|
||||
const int64_t nheads = final_states.size(1);
|
||||
const int64_t headdim = final_states.size(2);
|
||||
const int64_t dstate = final_states.size(3);
|
||||
const int64_t ngroups = B_in.size(1);
|
||||
|
||||
TORCH_CHECK(final_states.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
|
||||
TORCH_CHECK(out.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
|
||||
"raw data_ptr)");
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
|
||||
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
|
||||
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
|
||||
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
|
||||
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
|
||||
ngroups, headdim, dstate);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Fused CPU vector kernels for Mamba decode-step hotspots:
|
||||
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
|
||||
// - selective_state_update (SSM recurrence, single-step)
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cstdint>
|
||||
#include <algorithm>
|
||||
|
||||
namespace mamba_cpu {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update — templated for native BF16/FP32
|
||||
//
|
||||
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
|
||||
// Explicit strides are passed so the kernel writes directly into the
|
||||
// correct memory locations without making a contiguous copy of the full
|
||||
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
|
||||
//
|
||||
// stride_s_slot = state.stride(0) — between cache slots
|
||||
// stride_s_dim = state.stride(1) — between conv_dim channels
|
||||
// stride_s_state = state.stride(2) — between state elements
|
||||
//
|
||||
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t>
|
||||
inline void causal_conv1d_update_kernel(
|
||||
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
|
||||
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
|
||||
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
|
||||
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
|
||||
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
|
||||
int64_t width, int64_t state_len, bool do_silu) {
|
||||
#pragma omp parallel for
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
|
||||
if (cache_idx == pad_slot_id) continue;
|
||||
|
||||
for (int64_t t = 0; t < seqlen; ++t) {
|
||||
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
|
||||
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
|
||||
// Base of this slot in the (possibly non-contiguous) paged state
|
||||
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_b[d * seqlen]);
|
||||
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
|
||||
const scalar_t* w = weight_ptr + d * width;
|
||||
|
||||
// Accumulate in float32 for precision
|
||||
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
|
||||
for (int64_t k = 0; k < state_len; ++k) {
|
||||
acc += static_cast<float>(w[k]) *
|
||||
static_cast<float>(sd[k * stride_s_state]);
|
||||
}
|
||||
acc += static_cast<float>(w[state_len]) * x_val;
|
||||
|
||||
// Shift state left and append new input.
|
||||
// Use memmove when contiguous (stride==1); element loop otherwise.
|
||||
if (stride_s_state == 1) {
|
||||
if (state_len > 1)
|
||||
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
|
||||
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
|
||||
} else {
|
||||
for (int64_t k = 0; k < state_len - 1; ++k)
|
||||
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
|
||||
if (state_len > 0)
|
||||
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
|
||||
}
|
||||
|
||||
if (do_silu) {
|
||||
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
|
||||
: std::exp(acc) / (1.0f + std::exp(acc));
|
||||
acc *= sigmoid;
|
||||
}
|
||||
out_b[d * seqlen] = static_cast<scalar_t>(acc);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
//
|
||||
// Template parameters:
|
||||
// state_t - dtype of ssm_state cache (typically BFloat16)
|
||||
// input_t - dtype of x, B, C (typically BFloat16)
|
||||
// out_t - dtype of output tensor (typically BFloat16)
|
||||
// Write directly — no float32 intermediate buffer needed.
|
||||
//
|
||||
// A, D, dt_bias are accepted as const float* (they are always float32
|
||||
// model parameters in Mamba2). This eliminates the per-call float32→BF16
|
||||
// conversion and the .contiguous() materialisation of the broadcast-expand.
|
||||
//
|
||||
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
|
||||
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename state_t, typename input_t, typename out_t = float>
|
||||
inline void selective_state_update_kernel(
|
||||
state_t* __restrict__ state_ptr, int64_t stride_state_n,
|
||||
int64_t stride_state_h, int64_t stride_state_d,
|
||||
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
|
||||
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
|
||||
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
|
||||
// A: (nheads,) float32 — scalar per head
|
||||
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
|
||||
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
|
||||
// D: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ D_ptr,
|
||||
// z: same shape as x (optional)
|
||||
const input_t* __restrict__ z_ptr,
|
||||
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
|
||||
int64_t stride_out_n, int64_t stride_out_h,
|
||||
const int32_t* __restrict__ state_batch_indices,
|
||||
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
|
||||
const int32_t* __restrict__ num_accepted_tokens,
|
||||
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
|
||||
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
|
||||
using state_vec_t = vec_op::vec_t<state_t>;
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
int64_t nheads_per_group = nheads / ngroups;
|
||||
|
||||
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
|
||||
int64_t bos, seq_len;
|
||||
if (cu_seqlens != nullptr) {
|
||||
bos = cu_seqlens[seq_idx];
|
||||
seq_len = cu_seqlens[seq_idx + 1] - bos;
|
||||
} else {
|
||||
bos = seq_idx;
|
||||
seq_len = 1;
|
||||
}
|
||||
|
||||
int64_t state_read_idx = (state_batch_indices != nullptr)
|
||||
? state_batch_indices[seq_idx]
|
||||
: seq_idx;
|
||||
if (state_read_idx == null_block_id) continue;
|
||||
|
||||
int64_t state_write_idx = (num_accepted_tokens == nullptr)
|
||||
? ((dst_state_batch_indices != nullptr)
|
||||
? dst_state_batch_indices[seq_idx]
|
||||
: state_read_idx)
|
||||
: -1;
|
||||
|
||||
state_t* s = state_ptr + state_read_idx * stride_state_n;
|
||||
|
||||
for (int64_t t = 0; t < seq_len; ++t) {
|
||||
int64_t token_idx = bos + t;
|
||||
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
|
||||
// dt: (N, nheads) — one float per head per token
|
||||
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
|
||||
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
|
||||
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
|
||||
out_t* out_tok = out_ptr + token_idx * stride_out_n;
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
int64_t g = h / nheads_per_group;
|
||||
const input_t* x_h = x_tok + h * stride_x_h;
|
||||
const input_t* B_g = B_tok + g * stride_BC_g;
|
||||
const input_t* C_g = C_tok + g * stride_BC_g;
|
||||
out_t* out_h = out_tok + h * stride_out_h;
|
||||
state_t* s_h = s + h * stride_state_h;
|
||||
|
||||
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
|
||||
float dt_val = dt_tok[h];
|
||||
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
|
||||
if (dt_softplus) {
|
||||
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
|
||||
}
|
||||
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
const input_t* z_h =
|
||||
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
|
||||
: nullptr;
|
||||
|
||||
vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
|
||||
// and broadcast. This saves 7 redundant std::exp() calls that
|
||||
// FP32Vec8::exp() would otherwise make on the broadcast vector.
|
||||
const float dA_scalar = std::exp(A_val * dt_val);
|
||||
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_h[d]);
|
||||
|
||||
vec_op::FP32Vec8 out_vec(0.0f);
|
||||
state_t* s_hd = s_h + d * stride_state_d;
|
||||
const input_t* B_g_base = B_g;
|
||||
const input_t* C_g_base = C_g;
|
||||
|
||||
vec_op::FP32Vec8 x_vec(x_val);
|
||||
// dBx = B * x * dt — same dA for all dstate (A is scalar)
|
||||
// s_new = s * dA + B * x * dt
|
||||
|
||||
int64_t n = 0;
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
|
||||
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
|
||||
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
|
||||
|
||||
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
|
||||
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
|
||||
|
||||
state_vec_t(s_new).save(s_hd + n);
|
||||
out_vec = out_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float out_val = out_vec.reduce_sum();
|
||||
for (; n < dstate; ++n) {
|
||||
// Reuse dA_scalar computed once per head — no exp() re-call
|
||||
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
|
||||
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
|
||||
s_hd[n] = static_cast<state_t>(s_new);
|
||||
out_val += s_new * static_cast<float>(C_g[n]);
|
||||
}
|
||||
|
||||
if (D_ptr != nullptr) out_val += x_val * D_val;
|
||||
if (z_h != nullptr) {
|
||||
float z_val = static_cast<float>(z_h[d]);
|
||||
float sigmoid = (z_val >= 0)
|
||||
? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
out_val *= z_val * sigmoid;
|
||||
}
|
||||
out_h[d] = static_cast<out_t>(out_val);
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens != nullptr &&
|
||||
dst_state_batch_indices != nullptr) {
|
||||
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
|
||||
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
|
||||
state_write_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd
|
||||
//
|
||||
// Prefill SSM recurrence for Mamba2 / SSD models.
|
||||
//
|
||||
// Key difference from selective_state_update_kernel (decode path):
|
||||
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
|
||||
// Each thread owns a (batch, head) slice and runs the entire token
|
||||
// sequence without any per-token OpenMP synchronisation overhead.
|
||||
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
|
||||
//
|
||||
// `dt` arrives already processed (float32, after bias + softplus + clamp)
|
||||
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
|
||||
//
|
||||
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
|
||||
// tensor, pre-initialised by the caller (zero or from initial_states).
|
||||
// Each (b, h) slice is private to exactly one thread via collapse(2), so
|
||||
// there are no write conflicts.
|
||||
//
|
||||
// D is treated as a scalar per head ([nheads] float32).
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename input_t>
|
||||
inline void mamba_chunk_scan_fwd_kernel(
|
||||
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
|
||||
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
|
||||
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
|
||||
const float* __restrict__ A_ptr, // [nheads] f32
|
||||
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
|
||||
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
|
||||
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
|
||||
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
|
||||
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
|
||||
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
|
||||
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
|
||||
int64_t dstate) {
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
const int64_t nheads_per_group = nheads / ngroups;
|
||||
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
|
||||
// guarantee)
|
||||
const int64_t stride_s_b = nheads * headdim * dstate;
|
||||
const int64_t stride_s_h = headdim * dstate;
|
||||
// stride_s_d = dstate, stride_s_n = 1
|
||||
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
const int64_t seq_start = cu_seqlens[b];
|
||||
const int64_t seq_end = cu_seqlens[b + 1];
|
||||
const int64_t g = h / nheads_per_group;
|
||||
|
||||
const float A_val = A_ptr[h];
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
|
||||
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
|
||||
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
|
||||
|
||||
for (int64_t t = seq_start; t < seq_end; ++t) {
|
||||
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
|
||||
const float* dt_h = dt_ptr + t * nheads + h;
|
||||
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* z_h = (z_ptr != nullptr)
|
||||
? z_ptr + t * nheads * headdim + h * headdim
|
||||
: nullptr;
|
||||
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
|
||||
|
||||
const float dt_val = *dt_h;
|
||||
const float dA_val = std::exp(A_val * dt_val);
|
||||
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
|
||||
const vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
|
||||
for (int64_t d = 0; d < headdim; ++d) {
|
||||
const float x_val = static_cast<float>(x_h[d]);
|
||||
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
|
||||
|
||||
// Vectorised SSM update + readout over dstate:
|
||||
// s_new = s * dA + x * dt * B
|
||||
// y += s_new * C
|
||||
int64_t n = 0;
|
||||
vec_op::FP32Vec8 y_vec(0.0f);
|
||||
const vec_op::FP32Vec8 x_vec(x_val);
|
||||
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
|
||||
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
|
||||
const vec_op::FP32Vec8 s_v(s_bhd + n);
|
||||
|
||||
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
|
||||
s_new.save(s_bhd + n);
|
||||
y_vec = y_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float y_val = y_vec.reduce_sum();
|
||||
|
||||
// Scalar tail for remaining dstate elements
|
||||
for (; n < dstate; ++n) {
|
||||
const float B_n = static_cast<float>(B_g[n]);
|
||||
const float C_n = static_cast<float>(C_g[n]);
|
||||
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
|
||||
s_bhd[n] = s_new;
|
||||
y_val += s_new * C_n;
|
||||
}
|
||||
|
||||
// D skip connection (scalar per head)
|
||||
if (D_ptr != nullptr) y_val += x_val * D_val;
|
||||
|
||||
// z gating: out = y * z * sigmoid(z) (SiLU)
|
||||
if (z_h != nullptr) {
|
||||
const float z_val = static_cast<float>(z_h[d]);
|
||||
const float sigmoid =
|
||||
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
y_val *= z_val * sigmoid;
|
||||
}
|
||||
|
||||
out_h[d] = static_cast<input_t>(y_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mamba_cpu
|
||||
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
|
||||
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens);
|
||||
|
||||
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
|
||||
const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const at::Tensor& cu_seqlens);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
@@ -570,7 +596,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
#endif
|
||||
|
||||
// fused moe
|
||||
#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
|
||||
#if defined(__AVX512F__) || \
|
||||
(defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
|
||||
ops.def(
|
||||
"prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
|
||||
"-> ()");
|
||||
@@ -581,7 +608,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"bool skip_weighted, "
|
||||
"str act, str isa) -> ()");
|
||||
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
|
||||
#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
|
||||
#endif
|
||||
ops.def(
|
||||
"mla_decode_kvcache("
|
||||
" Tensor! out, Tensor query, Tensor kv_cache,"
|
||||
@@ -594,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
// Mamba CPU kernels
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu_vec("
|
||||
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
|
||||
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
|
||||
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
|
||||
&causal_conv1d_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"selective_state_update_cpu("
|
||||
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
|
||||
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
|
||||
"SymInt null_block_id, Tensor(a13!) out, "
|
||||
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
|
||||
&selective_state_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"mamba_chunk_scan_fwd_cpu("
|
||||
"Tensor(a0!) out, Tensor(a1!) final_states, "
|
||||
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
|
||||
&mamba_chunk_scan_fwd_cpu_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
|
||||
@@ -1174,7 +1174,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
const int32_t num_reqs, const int32_t block_size,
|
||||
const int32_t total_tokens, const int64_t block_table_stride,
|
||||
const int64_t cache_block_stride, const int64_t cache_entry_stride,
|
||||
const int64_t dst_entry_stride) {
|
||||
const int64_t dst_entry_stride,
|
||||
const int32_t* __restrict__ seq_starts) { // Optional source offsets
|
||||
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
|
||||
if (flat_warp_id >= total_tokens) return;
|
||||
const int lane_id = threadIdx.x & 31;
|
||||
@@ -1192,7 +1193,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
|
||||
// Compute physical token address via block table
|
||||
const int out_token_id = flat_warp_id;
|
||||
const int token_offset = out_token_id - workspace_starts[req_id];
|
||||
int token_offset = out_token_id - workspace_starts[req_id];
|
||||
if (seq_starts != nullptr) token_offset += seq_starts[req_id];
|
||||
const int cache_block_idx = token_offset / block_size;
|
||||
const int offset_in_block = token_offset % block_size;
|
||||
const int physical_block =
|
||||
@@ -1383,9 +1385,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
|
||||
torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::stable::Tensor const& seq_lens, // [BATCH]
|
||||
torch::stable::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size) {
|
||||
int64_t batch_size,
|
||||
std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
|
||||
torch::stable::accelerator::DeviceGuard device_guard(
|
||||
src_cache.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
@@ -1396,20 +1398,25 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
STD_TORCH_CHECK(
|
||||
block_table.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"block_table must be int32");
|
||||
STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"seq_lens must be int32");
|
||||
STD_TORCH_CHECK(
|
||||
workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"workspace_starts must be int32");
|
||||
if (seq_starts.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"seq_starts must be int32");
|
||||
}
|
||||
|
||||
STD_TORCH_CHECK(src_cache.device() == dst.device(),
|
||||
"src_cache and dst must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == block_table.device(),
|
||||
"src_cache and block_table must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),
|
||||
"src_cache and seq_lens must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
|
||||
"src_cache and workspace_starts must be on the same device");
|
||||
if (seq_starts.has_value()) {
|
||||
STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
|
||||
"src_cache and seq_starts must be on the same device");
|
||||
}
|
||||
auto dtype = src_cache.scalar_type();
|
||||
STD_TORCH_CHECK(
|
||||
dtype == torch::headeronly::ScalarType::Byte || // uint8
|
||||
@@ -1438,6 +1445,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
constexpr int warps_per_block = 8;
|
||||
const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
|
||||
const int block_size_threads = warps_per_block * 32; // 256 threads
|
||||
const int32_t* seq_starts_ptr =
|
||||
seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
|
||||
: nullptr;
|
||||
|
||||
vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
|
||||
stream>>>(
|
||||
@@ -1446,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
workspace_starts.const_data_ptr<int32_t>(),
|
||||
static_cast<int32_t>(batch_size), block_size, total_tokens,
|
||||
block_table_stride, cache_block_stride, cache_entry_stride,
|
||||
dst_entry_stride);
|
||||
dst_entry_stride, seq_starts_ptr);
|
||||
}
|
||||
|
||||
// Macro to dispatch the kernel based on the data type.
|
||||
|
||||
@@ -21,6 +21,21 @@
|
||||
#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \
|
||||
THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
|
||||
|
||||
// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids)
|
||||
// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'.
|
||||
#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \
|
||||
THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__)
|
||||
|
||||
#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \
|
||||
VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \
|
||||
__VA_ARGS__) \
|
||||
VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \
|
||||
__VA_ARGS__)
|
||||
|
||||
#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \
|
||||
THO_DISPATCH_SWITCH(TYPE, NAME, \
|
||||
VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__))
|
||||
|
||||
#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
|
||||
THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
|
||||
|
||||
@@ -360,12 +360,24 @@ __global__ void count_and_sort_expert_tokens_kernel(
|
||||
template <typename scalar_t>
|
||||
constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t);
|
||||
|
||||
template <typename scalar_t, int TOPK>
|
||||
template <typename idx_t>
|
||||
__device__ __forceinline__ bool moe_sum_pad_aware_skip(
|
||||
const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
|
||||
int64_t idx) {
|
||||
int64_t expert_id = static_cast<int64_t>(topk_ids[idx]);
|
||||
if (expert_id < 0) return true;
|
||||
if (expert_map != nullptr && expert_map[expert_id] < 0) return true;
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename idx_t, int TOPK, bool PAD_AWARE>
|
||||
__global__ void moe_sum_vec_kernel(
|
||||
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
|
||||
const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
|
||||
const int64_t num_tokens, const int d, const int64_t stride_token,
|
||||
const int64_t stride_topk) {
|
||||
const int64_t stride_topk, const idx_t* __restrict__ topk_ids,
|
||||
const int32_t* __restrict__ expert_map, const int64_t stride_tk_token,
|
||||
const int64_t stride_tk_k) {
|
||||
using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>; // 16-byte pack
|
||||
constexpr int VEC = MOE_SUM_VEC<scalar_t>;
|
||||
const int64_t n_vec = d / VEC;
|
||||
@@ -375,6 +387,10 @@ __global__ void moe_sum_vec_kernel(
|
||||
const int64_t token = i / n_vec;
|
||||
const int64_t v = i % n_vec;
|
||||
const scalar_t* in_tok = input + token * stride_token + v * VEC;
|
||||
const idx_t* tk_tok = nullptr;
|
||||
if constexpr (PAD_AWARE) {
|
||||
tk_tok = topk_ids + token * stride_tk_token;
|
||||
}
|
||||
|
||||
float acc[VEC];
|
||||
#pragma unroll
|
||||
@@ -382,6 +398,11 @@ __global__ void moe_sum_vec_kernel(
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < TOPK; ++k) {
|
||||
if constexpr (PAD_AWARE) {
|
||||
if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
|
||||
@@ -394,13 +415,16 @@ __global__ void moe_sum_vec_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
// Runtime-topk variant of the above.
|
||||
template <typename scalar_t>
|
||||
// Runtime-topk variant of the above, for topk values outside the templated
|
||||
// set.
|
||||
template <typename scalar_t, typename idx_t, bool PAD_AWARE>
|
||||
__global__ void moe_sum_vec_dynamic_kernel(
|
||||
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
|
||||
const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
|
||||
const int64_t num_tokens, const int d, const int topk,
|
||||
const int64_t stride_token, const int64_t stride_topk) {
|
||||
const int64_t stride_token, const int64_t stride_topk,
|
||||
const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
|
||||
const int64_t stride_tk_token, const int64_t stride_tk_k) {
|
||||
using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>;
|
||||
constexpr int VEC = MOE_SUM_VEC<scalar_t>;
|
||||
const int64_t n_vec = d / VEC;
|
||||
@@ -410,12 +434,21 @@ __global__ void moe_sum_vec_dynamic_kernel(
|
||||
const int64_t token = i / n_vec;
|
||||
const int64_t v = i % n_vec;
|
||||
const scalar_t* in_tok = input + token * stride_token + v * VEC;
|
||||
const idx_t* tk_tok = nullptr;
|
||||
if constexpr (PAD_AWARE) {
|
||||
tk_tok = topk_ids + token * stride_tk_token;
|
||||
}
|
||||
|
||||
float acc[VEC];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
|
||||
|
||||
for (int k = 0; k < topk; ++k) {
|
||||
if constexpr (PAD_AWARE) {
|
||||
if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
|
||||
@@ -430,17 +463,28 @@ __global__ void moe_sum_vec_dynamic_kernel(
|
||||
|
||||
// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims
|
||||
// (including a non-contiguous hidden stride) via per-element strided reads.
|
||||
template <typename scalar_t>
|
||||
template <typename scalar_t, typename idx_t, bool PAD_AWARE>
|
||||
__global__ void moe_sum_scalar_kernel(
|
||||
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
|
||||
const scalar_t* __restrict__ input, // [num_tokens, topk, d]
|
||||
const int d, const int topk, const int64_t stride_token,
|
||||
const int64_t stride_topk, const int64_t stride_hidden) {
|
||||
const int64_t stride_topk, const int64_t stride_hidden,
|
||||
const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
|
||||
const int64_t stride_tk_token, const int64_t stride_tk_k) {
|
||||
const int64_t token_idx = blockIdx.x;
|
||||
const scalar_t* in_tok = input + token_idx * stride_token;
|
||||
const idx_t* tk_tok = nullptr;
|
||||
if constexpr (PAD_AWARE) {
|
||||
tk_tok = topk_ids + token_idx * stride_tk_token;
|
||||
}
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
float x = 0.f;
|
||||
for (int k = 0; k < topk; ++k) {
|
||||
if constexpr (PAD_AWARE) {
|
||||
if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
x += static_cast<float>(
|
||||
VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden]));
|
||||
}
|
||||
@@ -711,8 +755,9 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
|
||||
}
|
||||
|
||||
void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
|
||||
torch::stable::Tensor& output) // [num_tokens, hidden_size]
|
||||
{
|
||||
torch::stable::Tensor& output, // [num_tokens, hidden_size]
|
||||
std::optional<torch::stable::Tensor> topk_ids,
|
||||
std::optional<torch::stable::Tensor> expert_map) {
|
||||
// Output is dense and written in place, so it must be contiguous. The input
|
||||
// is read by its strides (no copy); only the hidden dim needs to be
|
||||
// contiguous to take the vectorized path.
|
||||
@@ -731,10 +776,102 @@ void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(output.get_device_index());
|
||||
|
||||
#define LAUNCH_MOE_SUM_VEC(TOPK) \
|
||||
vllm::moe::moe_sum_vec_kernel<scalar_t, TOPK> \
|
||||
<<<grid, dim3(block), 0, stream>>>( \
|
||||
out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk)
|
||||
if (topk_ids.has_value()) {
|
||||
// Pad-aware reduce path
|
||||
const torch::stable::Tensor& tk = topk_ids.value();
|
||||
STD_TORCH_CHECK(tk.size(0) == num_tokens && tk.size(1) == topk,
|
||||
"moe_sum: topk_ids must have shape [num_tokens, topk]");
|
||||
const int64_t stride_tk_token = tk.stride(0);
|
||||
const int64_t stride_tk_k = tk.stride(1);
|
||||
|
||||
const int32_t* expert_map_ptr = nullptr;
|
||||
if (expert_map.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
expert_map->scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"moe_sum: expert_map must be int32");
|
||||
expert_map_ptr =
|
||||
reinterpret_cast<const int32_t*>(expert_map->const_data_ptr());
|
||||
}
|
||||
|
||||
#define LAUNCH_MOE_SUM_PAD_AWARE_VEC(TOPK) \
|
||||
vllm::moe::moe_sum_vec_kernel<scalar_t, idx_t, TOPK, true> \
|
||||
<<<grid, dim3(block), 0, stream>>>( \
|
||||
out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk, \
|
||||
topk_ids_ptr, expert_map_ptr, stride_tk_token, stride_tk_k)
|
||||
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "moe_sum_pad_aware", [&] {
|
||||
constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
|
||||
constexpr int WIDTH = VEC * sizeof(scalar_t);
|
||||
auto* out_ptr =
|
||||
reinterpret_cast<scalar_t*>(output.mutable_data_ptr());
|
||||
auto* in_ptr =
|
||||
reinterpret_cast<const scalar_t*>(input.const_data_ptr());
|
||||
|
||||
const bool can_vec =
|
||||
(stride_hidden == 1) && (hidden_size % VEC == 0) &&
|
||||
(stride_token % VEC == 0) && (stride_topk % VEC == 0) &&
|
||||
(reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) &&
|
||||
(reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0);
|
||||
|
||||
VLLM_STABLE_DISPATCH_IDX_TYPES(
|
||||
tk.scalar_type(), "moe_sum_pad_aware_idx", [&] {
|
||||
auto* topk_ids_ptr =
|
||||
reinterpret_cast<const idx_t*>(tk.const_data_ptr());
|
||||
if (can_vec) {
|
||||
const int64_t n_vec = hidden_size / VEC;
|
||||
const int64_t total = num_tokens * n_vec;
|
||||
const int block = 256;
|
||||
const dim3 grid(
|
||||
std::min<int64_t>((total + block - 1) / block, 65535));
|
||||
switch (topk) {
|
||||
case 1:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(1);
|
||||
break;
|
||||
case 2:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(2);
|
||||
break;
|
||||
case 4:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(4);
|
||||
break;
|
||||
case 6:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(6);
|
||||
break;
|
||||
case 8:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(8);
|
||||
break;
|
||||
case 9:
|
||||
LAUNCH_MOE_SUM_PAD_AWARE_VEC(9);
|
||||
break;
|
||||
default:
|
||||
vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, idx_t,
|
||||
true>
|
||||
<<<grid, dim3(block), 0, stream>>>(
|
||||
out_ptr, in_ptr, num_tokens, hidden_size, topk,
|
||||
stride_token, stride_topk, topk_ids_ptr,
|
||||
expert_map_ptr, stride_tk_token, stride_tk_k);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
dim3 grid(num_tokens);
|
||||
dim3 block(std::min(hidden_size, 1024));
|
||||
vllm::moe::moe_sum_scalar_kernel<scalar_t, idx_t, true>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out_ptr, in_ptr, hidden_size, topk, stride_token,
|
||||
stride_topk, stride_hidden, topk_ids_ptr,
|
||||
expert_map_ptr, stride_tk_token, stride_tk_k);
|
||||
}
|
||||
});
|
||||
});
|
||||
#undef LAUNCH_MOE_SUM_PAD_AWARE_VEC
|
||||
return;
|
||||
}
|
||||
|
||||
#define LAUNCH_MOE_SUM_VEC(TOPK) \
|
||||
vllm::moe::moe_sum_vec_kernel<scalar_t, int32_t, TOPK, false> \
|
||||
<<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens, \
|
||||
hidden_size, stride_token, \
|
||||
stride_topk, nullptr, nullptr, 0, 0)
|
||||
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] {
|
||||
constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
|
||||
@@ -774,18 +911,19 @@ void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
|
||||
LAUNCH_MOE_SUM_VEC(9);
|
||||
break;
|
||||
default:
|
||||
vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t>
|
||||
<<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens,
|
||||
hidden_size, topk,
|
||||
stride_token, stride_topk);
|
||||
vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, int32_t, false>
|
||||
<<<grid, dim3(block), 0, stream>>>(
|
||||
out_ptr, in_ptr, num_tokens, hidden_size, topk, stride_token,
|
||||
stride_topk, nullptr, nullptr, 0, 0);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
dim3 grid(num_tokens);
|
||||
dim3 block(std::min(hidden_size, 1024));
|
||||
vllm::moe::moe_sum_scalar_kernel<scalar_t><<<grid, block, 0, stream>>>(
|
||||
out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk,
|
||||
stride_hidden);
|
||||
vllm::moe::moe_sum_scalar_kernel<scalar_t, int32_t, false>
|
||||
<<<grid, block, 0, stream>>>(out_ptr, in_ptr, hidden_size, topk,
|
||||
stride_token, stride_topk, stride_hidden,
|
||||
nullptr, nullptr, 0, 0);
|
||||
}
|
||||
});
|
||||
#undef LAUNCH_MOE_SUM_VEC
|
||||
|
||||
@@ -27,7 +27,9 @@ void topk_softplus_sqrt(
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid);
|
||||
|
||||
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
|
||||
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,
|
||||
|
||||
@@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162;
|
||||
namespace vllm {
|
||||
namespace moe {
|
||||
|
||||
template <typename HashIndType>
|
||||
__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr,
|
||||
int64_t offset) {
|
||||
return static_cast<int64_t>(ptr[offset]);
|
||||
}
|
||||
|
||||
/// Aligned array type
|
||||
template <typename T,
|
||||
/// Number of elements in the array
|
||||
@@ -65,6 +71,73 @@ __device__ __forceinline__ float toFloat(T value) {
|
||||
}
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Adapted from:
|
||||
// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
|
||||
template <typename OutIndType, typename HashIndType>
|
||||
__launch_bounds__(128) __global__
|
||||
void dsv4HashTopkSoftplusSqrt(const float* input, float* output,
|
||||
OutIndType* indices, int num_rows,
|
||||
int num_experts, float routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid) {
|
||||
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
if (warp >= num_rows) return;
|
||||
const int64_t token_id = load_index_as_int64(input_ids, warp);
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
int expert = 0;
|
||||
float weight = 0.f;
|
||||
if (lane < 6) {
|
||||
// only load and calculate for 6 experts
|
||||
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
|
||||
const float x = input[warp * num_experts + expert];
|
||||
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
|
||||
}
|
||||
float weight_sum = weight;
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1) {
|
||||
// sum in warp
|
||||
weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask);
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
if (lane < 6) {
|
||||
const int offset = warp * 6 + lane;
|
||||
output[offset] =
|
||||
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
|
||||
indices[offset] = static_cast<OutIndType>(expert);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutIndType, typename HashIndType>
|
||||
void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
int num_rows, int num_experts,
|
||||
double routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
if (num_rows == 0) return;
|
||||
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
|
||||
cudaLaunchConfig_t config = {};
|
||||
config.gridDim = (num_rows + 3) / 4;
|
||||
config.blockDim = 128;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attr;
|
||||
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attr.val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = &attr;
|
||||
config.numAttrs = 1;
|
||||
const float scale = static_cast<float>(routed_scaling_factor);
|
||||
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
|
||||
num_experts, scale, input_ids, tid2eid);
|
||||
}
|
||||
#endif
|
||||
|
||||
// ====================== TopK softplus_sqrt things
|
||||
// ===============================
|
||||
|
||||
@@ -86,14 +159,14 @@ __device__ __forceinline__ float toFloat(T value) {
|
||||
|
||||
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
|
||||
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
|
||||
typename InputType = float>
|
||||
typename HashIndType, typename InputType = float>
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGatingSoftplusSqrt(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const IndType* input_ids, const IndType* tid2eid) {
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -240,8 +313,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
|
||||
// Hash MoE path: indices are predetermined from lookup table
|
||||
if constexpr (USE_HASH) {
|
||||
const IndType token_id = input_ids[thread_row];
|
||||
const IndType* expert_indices_for_token = tid2eid + token_id * k;
|
||||
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
|
||||
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
@@ -252,7 +325,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int expert = static_cast<int>(
|
||||
load_index_as_int64(tid2eid, token_expert_offset + k_idx));
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
@@ -261,7 +335,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = expert;
|
||||
indices[idx] = static_cast<IndType>(expert);
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
@@ -285,7 +359,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int expert = static_cast<int>(
|
||||
load_index_as_int64(tid2eid, token_expert_offset + k_idx));
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
@@ -461,14 +536,15 @@ struct TopkConstants {
|
||||
}
|
||||
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
|
||||
int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
|
||||
int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType,
|
||||
typename InputType>
|
||||
void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
IndType* indices, int* source_row, const int num_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
const bool use_hash, const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
@@ -481,7 +557,8 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
DISPATCH_HASH(use_hash, USE_HASH, {
|
||||
auto* kernel =
|
||||
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
|
||||
WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
|
||||
WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType,
|
||||
InputType>;
|
||||
#ifndef USE_ROCM
|
||||
cudaLaunchConfig_t config = {};
|
||||
config.gridDim = num_blocks;
|
||||
@@ -538,13 +615,25 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename IndType, typename InputType>
|
||||
template <typename IndType, typename InputType, typename HashIndType = IndType>
|
||||
void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
|
||||
int* token_expert_indices, const int num_tokens, const int num_experts,
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash, const IndType* input_ids,
|
||||
const IndType* tid2eid, cudaStream_t stream) {
|
||||
const float* correction_bias, const bool use_hash,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
#ifndef USE_ROCM
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (use_hash && topk == 6 && renormalize &&
|
||||
(num_experts == 256 || num_experts == 384)) {
|
||||
launchDsv4HashTopk<IndType, HashIndType>(
|
||||
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
|
||||
routed_scaling_factor, input_ids, tid2eid, stream);
|
||||
return;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
|
||||
@@ -644,57 +733,55 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().const_data_ptr<float>();
|
||||
}
|
||||
bool use_hash = false;
|
||||
if (tid2eid.has_value()) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(),
|
||||
"input_ids is required for hash MoE");
|
||||
use_hash = true;
|
||||
}
|
||||
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
|
||||
const int* input_ids_ptr = nullptr;
|
||||
const int* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().const_data_ptr<int>();
|
||||
tid2eid_ptr = tid2eid.value().const_data_ptr<int>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices.mutable_data_ptr<int>(),
|
||||
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
|
||||
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
|
||||
input_ids_ptr, tid2eid_ptr, stream);
|
||||
auto launch = [&](auto* topk_indices_ptr) {
|
||||
using OutIndType =
|
||||
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
|
||||
if (tid2eid.has_value()) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(),
|
||||
"input_ids is required for hash MoE");
|
||||
STD_TORCH_CHECK(
|
||||
input_ids.value().scalar_type() == tid2eid.value().scalar_type(),
|
||||
"input_ids and tid2eid must have the same dtype");
|
||||
if (tid2eid.value().scalar_type() ==
|
||||
torch::headeronly::ScalarType::Long) {
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
|
||||
int64_t>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream);
|
||||
} else {
|
||||
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
|
||||
torch::headeronly::ScalarType::Int);
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
|
||||
int>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
|
||||
tid2eid.value().const_data_ptr<int>(), stream);
|
||||
}
|
||||
} else {
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
|
||||
static_cast<const OutIndType*>(nullptr), stream);
|
||||
}
|
||||
};
|
||||
|
||||
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
|
||||
launch(topk_indices.mutable_data_ptr<int>());
|
||||
} else if (topk_indices.scalar_type() ==
|
||||
torch::headeronly::ScalarType::UInt32) {
|
||||
const uint32_t* input_ids_ptr = nullptr;
|
||||
const uint32_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().const_data_ptr<uint32_t>();
|
||||
tid2eid_ptr = tid2eid.value().const_data_ptr<uint32_t>();
|
||||
}
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices.mutable_data_ptr<uint32_t>(),
|
||||
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
|
||||
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
|
||||
input_ids_ptr, tid2eid_ptr, stream);
|
||||
launch(topk_indices.mutable_data_ptr<uint32_t>());
|
||||
} else {
|
||||
STD_TORCH_CHECK(topk_indices.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Long);
|
||||
|
||||
const int64_t* input_ids_ptr = nullptr;
|
||||
const int64_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().const_data_ptr<int64_t>();
|
||||
tid2eid_ptr = tid2eid.value().const_data_ptr<int64_t>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
|
||||
gating_output, topk_weights.mutable_data_ptr<float>(),
|
||||
topk_indices.mutable_data_ptr<int64_t>(),
|
||||
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
|
||||
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
|
||||
input_ids_ptr, tid2eid_ptr, stream);
|
||||
launch(topk_indices.mutable_data_ptr<int64_t>());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -738,4 +825,4 @@ void topk_softplus_sqrt(
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,8 +23,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts.
|
||||
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
|
||||
// from all selected experts. topk_ids/expert_map are optional and, when
|
||||
// both given, enable pad-aware reduce that skips (token, expert)
|
||||
// slots that were never actually computed (unrouted, or routed to an
|
||||
// expert not owned by this rank under expert parallelism).
|
||||
m.def(
|
||||
"moe_sum(Tensor input, Tensor! output, Tensor? topk_ids=None, "
|
||||
"Tensor? expert_map=None) -> ()");
|
||||
|
||||
// Aligning the number of tokens to be processed by each expert such
|
||||
// that it is divisible by the block size.
|
||||
|
||||
@@ -527,9 +527,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
// 656]
|
||||
torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::stable::Tensor const& seq_lens, // [BATCH]
|
||||
torch::stable::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size);
|
||||
int64_t batch_size,
|
||||
std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
|
||||
|
||||
// Indexer K quantization and cache function
|
||||
void indexer_k_quant_and_cache(
|
||||
|
||||
@@ -847,8 +847,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
|
||||
|
||||
ops.def(
|
||||
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
|
||||
"Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
|
||||
"batch_size) -> ()");
|
||||
"Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
|
||||
"seq_starts) -> ()");
|
||||
|
||||
ops.def(
|
||||
"indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
|
||||
|
||||
+1
-36
@@ -248,10 +248,6 @@ COPY requirements/common.txt requirements/common.txt
|
||||
COPY requirements/cuda.txt requirements/cuda.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. uv can extract them in either order,
|
||||
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
|
||||
@@ -268,13 +264,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
else \
|
||||
uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi \
|
||||
&& if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi
|
||||
|
||||
# Track PyTorch lib versions used during build and match in downstream instances.
|
||||
@@ -792,10 +781,6 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
|
||||
ARG PYTORCH_CUDA_INDEX_BASE_URL
|
||||
COPY requirements/common.txt /tmp/common.txt
|
||||
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. uv can extract them in either order,
|
||||
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
|
||||
@@ -803,19 +788,12 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
fi && \
|
||||
uv pip install --system -r /tmp/requirements-cuda.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --system --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi && \
|
||||
rm /tmp/requirements-cuda.txt /tmp/common.txt
|
||||
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
# https://docs.flashinfer.ai/installation.html
|
||||
# From versions.json: .flashinfer.version
|
||||
ARG FLASHINFER_VERSION=0.6.13
|
||||
ARG FLASHINFER_VERSION=0.6.14
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
@@ -908,19 +886,6 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. Force -libs-cu13 last after runtime
|
||||
# dependency installs so uv cannot leave base files behind.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --system --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi
|
||||
|
||||
# CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
|
||||
# return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
|
||||
# consistently from the host (see https://github.com/vllm-project/vllm/issues/18859).
|
||||
|
||||
@@ -674,6 +674,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
|
||||
libibverbs1 \
|
||||
ibverbs-providers \
|
||||
ibverbs-utils \
|
||||
util-linux \
|
||||
unzip \
|
||||
pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
|
||||
libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
|
||||
|
||||
@@ -68,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.13"
|
||||
"default": "0.6.14"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -5,7 +5,7 @@ vLLM uses the following environment variables to configure the system:
|
||||
!!! warning
|
||||
Please note that `VLLM_PORT` and `VLLM_HOST_IP` set the port and ip for vLLM's **internal usage**. It is not the port and ip for the API server. If you use `--host $VLLM_HOST_IP` and `--port $VLLM_PORT` to start the API server, it will not work.
|
||||
|
||||
All environment variables used by vLLM are prefixed with `VLLM_`. **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).
|
||||
Most vLLM-specific environment variables are prefixed with `VLLM_` (a handful of standard names — for example `CUDA_VISIBLE_DEVICES`, `MAX_JOBS`, `S3_ACCESS_KEY_ID`/`S3_SECRET_ACCESS_KEY`/`S3_ENDPOINT_URL`, `DO_NOT_TRACK`, `NO_COLOR` — are also read directly when set). **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).
|
||||
|
||||
```python
|
||||
--8<-- "vllm/envs.py:env-vars-definition"
|
||||
|
||||
@@ -6,7 +6,11 @@ vLLM maintains a per-commit wheel repository (commonly referred to as "nightly")
|
||||
|
||||
### Wheel Building
|
||||
|
||||
Wheels are built in the `Release` pipeline (`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch, with multiple variants:
|
||||
Wheels are built in the `Release` pipeline
|
||||
(`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch.
|
||||
Regular builds produce the CUDA 13.0 wheels for x86_64 and aarch64. Additional
|
||||
wheel variants and ROCm builds can be unblocked on demand and run automatically
|
||||
when `NIGHTLY=1`:
|
||||
|
||||
- **Backend variants**: `cpu` and `cuXXX` (e.g., `cu129`, `cu130`).
|
||||
- **Architecture variants**: `x86_64` and `aarch64`.
|
||||
|
||||
@@ -217,7 +217,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
image: vllm/vllm-openai:latest
|
||||
command: ["/bin/sh", "-c"]
|
||||
args: [
|
||||
"vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024"
|
||||
"vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024"
|
||||
]
|
||||
env:
|
||||
- name: HF_TOKEN
|
||||
@@ -306,7 +306,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
- SYS_PTRACE
|
||||
command: ["/bin/sh", "-c"]
|
||||
args: [
|
||||
"vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024"
|
||||
"vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024"
|
||||
]
|
||||
env:
|
||||
- name: HF_TOKEN
|
||||
|
||||
@@ -205,7 +205,7 @@ hardware and configuration.
|
||||
|
||||
| Backend | Description | Dtypes | Compute Cap. | Notes |
|
||||
| ------- | ----------- | ------ | ------------ | ----- |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
|
||||
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
|
||||
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
|
||||
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
|
||||
|
||||
### SiLU+Mul + Quantization (`fuse_act_quant`)
|
||||
|
||||
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
|
||||
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
|
||||
### RMSNorm + Padding (`fuse_act_padding`)
|
||||
|
||||
|
||||
@@ -101,7 +101,7 @@ The heartbeat is deferred to the next step once the handshake completes --- the
|
||||
|
||||
## Bidirectional KV Transfer
|
||||
|
||||
For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire and P recomputes. Future work may extend a symmetric heartbeat mechanism to this case.
|
||||
For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire. D communicates back the expiry time so P can know when blocks are expired and recompute. Because the deadline is a `perf_counter` value produced on D and the two engines run in separate processes (with unrelated clocks), P estimates the clock offset to D from the handshake round-trip and applies it before comparing the deadline against its own `perf_counter`. Future work may extend a symmetric heartbeat mechanism to this case.
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
|
||||
@@ -107,6 +107,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
- **Phi series**: `microsoft/Phi-3.5-mini-instruct`
|
||||
|
||||
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
|
||||
|
||||
|
||||
@@ -68,13 +68,14 @@ vllm serve <model> \
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
|
||||
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
|
||||
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
|
||||
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
|
||||
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
@@ -83,9 +84,11 @@ Each entry in `secondary_tiers` is a dict with a required `type` field plus tier
|
||||
|
||||
The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`.
|
||||
|
||||
Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the tier's storage location relative to the publishing vLLM instance. `LOCAL` marks storage local to that instance, while `REMOTE` marks storage that is not local to it. When the setting is omitted, locality is unspecified. vLLM does not infer it from the tier type, so an OBJ tier is not implicitly `REMOTE`. A KV event includes `locality` only when the tier explicitly configures it. This metadata describes the tier property without implying that a consumer can already route requests to its blocks.
|
||||
|
||||
### Filesystem (FS)
|
||||
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a filesystem directory.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
@@ -94,6 +97,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage
|
||||
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
|
||||
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
|
||||
| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
|
||||
| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured. |
|
||||
|
||||
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
|
||||
|
||||
@@ -134,6 +138,7 @@ The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object
|
||||
| `prefix` | no | `""` | Key prefix prepended to all object keys. |
|
||||
| `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. |
|
||||
| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
|
||||
| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured; OBJ does not imply `REMOTE`. |
|
||||
|
||||
`store_config` fields:
|
||||
|
||||
@@ -152,21 +157,30 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti
|
||||
|
||||
The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
|
||||
|
||||
PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
| `type` | yes | — | Must be `p2p`. |
|
||||
| `host` | no | `0.0.0.0` | Address the control socket binds to. |
|
||||
| `port` | no | `7777` | Port for the control socket. Must be reachable from peers. |
|
||||
| `host` | no | `$VLLM_P2P_SIDE_CHANNEL_HOST` (`localhost`) | Address the control socket binds to, used verbatim as the identity peers dial back. When omitted, resolves from the env var below. The `localhost` default binds loopback only — for cross-host P2P you **must** set it to the node's routable IP (see below). |
|
||||
| `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. |
|
||||
| `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
|
||||
| `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
|
||||
|
||||
The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
|
||||
|
||||
#### Environment Variables
|
||||
|
||||
Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence.
|
||||
|
||||
- `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide.
|
||||
- `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision.
|
||||
|
||||
## Tuning Tips
|
||||
|
||||
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
|
||||
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
|
||||
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
|
||||
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
|
||||
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
|
||||
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block announce %}
|
||||
<p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/">Click here</a> to view docs for the latest stable release.</p>
|
||||
<p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/{{ page.url }}">Click here</a> to view docs for the latest stable release.</p>
|
||||
{% endblock %}
|
||||
|
||||
@@ -31,10 +31,8 @@
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
|
||||
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
|
||||
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
|
||||
|
||||
@@ -279,11 +279,66 @@ the pooler assigned to each task has the following attributes by default:
|
||||
| `embed` | `LAST` | ✅︎ | ❌ |
|
||||
| `classify` | `LAST` | ❌ | ✅︎ |
|
||||
|
||||
When loading [Sentence Transformers](https://huggingface.co/sentence-transformers) models,
|
||||
its Sentence Transformers configuration file (`modules.json`) takes priority over the model's defaults.
|
||||
#### Resolution precedence
|
||||
|
||||
You can further customize this via the `--pooler-config` option,
|
||||
which takes priority over both the model's and Sentence Transformers' defaults.
|
||||
The pooling method and `use_activation` are resolved per field. An explicitly
|
||||
set field in `--pooler-config` takes precedence over Sentence Transformers
|
||||
metadata, which in turn takes precedence over the model architecture or task
|
||||
default. Fields left unset continue through the chain independently.
|
||||
|
||||
The current `PoolerConfig` has no `normalize` or `activation` field.
|
||||
`use_activation` controls whether the task's constructed normalization or
|
||||
classification activation is applied.
|
||||
|
||||
| Field | Source precedence | How to override |
|
||||
| ----- | ----------------- | --------------- |
|
||||
| Pooling method (`pooling_type`) | `--pooler-config` > boolean `pooling_mode_*` fields in the Pooling module referenced by Sentence Transformers `modules.json` > architecture default (`LAST` for sequence pooling and `ALL` for token pooling unless the architecture overrides it) | Set `{"pooling_type": "CLS"}`, or set `seq_pooling_type` / `tok_pooling_type` explicitly. |
|
||||
| Embedding normalization (`use_activation`) | `--pooler-config` > Sentence Transformers modules (`true` when a Normalize module is present, otherwise `false`) > pooling-task default (`true`) when no Sentence Transformers Pooling module is found | Set `{"use_activation": false}` to return unnormalized embeddings. |
|
||||
| Classification activation function | Hugging Face `problem_type` > Sentence Transformers activation metadata > sigmoid or softmax selected from the label count | The function cannot be selected through `--pooler-config`; set `{"use_activation": false}` to return logits instead. |
|
||||
|
||||
Sentence Transformers configurations using the newer compact `pooling_mode`
|
||||
string are not currently parsed; see [issue #45995](https://github.com/vllm-project/vllm/issues/45995).
|
||||
|
||||
For converted models and predefined models using the standard DispatchPooler
|
||||
adapters, `embed` and `token_embed` construct an L2-normalization head, while
|
||||
`classify` and `token_classify` construct the selected classification activation.
|
||||
In both cases, `use_activation` controls whether that head is applied. Models with
|
||||
custom poolers can implement different behavior.
|
||||
|
||||
To inspect the resolved fields without loading model weights:
|
||||
|
||||
```python
|
||||
from vllm.config import ModelConfig, PoolerConfig
|
||||
from vllm.model_executor.layers.pooler.activations import get_act_fn
|
||||
|
||||
|
||||
def inspect(requested: PoolerConfig) -> None:
|
||||
model_config = ModelConfig(
|
||||
"intfloat/e5-small",
|
||||
runner="pooling",
|
||||
pooler_config=requested,
|
||||
)
|
||||
resolved = model_config.pooler_config
|
||||
assert resolved is not None
|
||||
print(
|
||||
{
|
||||
"seq_pooling_type": resolved.seq_pooling_type,
|
||||
"tok_pooling_type": resolved.tok_pooling_type,
|
||||
"use_activation": resolved.use_activation,
|
||||
"sequence_classification_activation": type(
|
||||
get_act_fn(model_config.hf_config)
|
||||
).__name__,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
inspect(PoolerConfig())
|
||||
inspect(PoolerConfig(pooling_type="CLS", use_activation=False))
|
||||
```
|
||||
|
||||
For `intfloat/e5-small`, the first result contains `MEAN`, `ALL`, and `True`.
|
||||
The second contains `CLS`, `ALL`, and `False`. Both report the classification
|
||||
activation that the standard sequence-classification adapter would construct.
|
||||
|
||||
## Removed Features
|
||||
|
||||
|
||||
@@ -47,6 +47,8 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
|
||||
| `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
|
||||
| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | |
|
||||
| `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
|
||||
| `RobertaForTokenClassification` | RoBERTa-based | `Jean-Baptiste/roberta-large-ner-english` | | |
|
||||
| `XLMRobertaForTokenClassification` | XLM-RoBERTa-based | `Davlan/xlm-roberta-base-ner-hrl` | | |
|
||||
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
|
||||
|
||||
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
|
||||
|
||||
@@ -514,6 +514,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ |
|
||||
| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | | ✅︎ |
|
||||
| `Cosmos3EdgeForConditionalGeneration` | Cosmos3-Edge (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Edge` | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -554,7 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
|
||||
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
|
||||
| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ |
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | ✅︎ | ✅︎ |
|
||||
| `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I<sup>+</sup> + V<sup>+</sup> | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
|
||||
@@ -220,8 +220,6 @@ For multi-node deployment, add these EPLB flags to each node's command. We recom
|
||||
|
||||
- Use simulator flags `VLLM_MOE_ROUTING_SIMULATION_STRATEGY=uniform_random` and `VLLM_RANDOMIZE_DP_DUMMY_INPUTS=1` so token routing is balanced across EP ranks.
|
||||
|
||||
- Increasing `VLLM_MOE_DP_CHUNK_SIZE` may increase throughput by increasing the maximum batch size for inter-rank token transfers. This may cause DeepEP to throw `assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2`, which can be fixed by increasing environment variable `NVSHMEM_QP_DEPTH`.
|
||||
|
||||
## Disaggregated Serving (Prefill/Decode Split)
|
||||
|
||||
For production deployments requiring strict SLA guarantees for time-to-first-token and inter-token latency, disaggregated serving allows independent scaling of prefill and decode operations.
|
||||
|
||||
@@ -137,8 +137,12 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
|
||||
|
||||
### Derenderer APIs
|
||||
|
||||
- `/v1/completions/derender` - Derenderer completion requests
|
||||
- `/v1/chat/completions/derender` - Derenderer chat completion requests
|
||||
For further details on derenderer APIs, please refer to [this page](derenderer.md).
|
||||
|
||||
- [Chat Completions Derender API](derenderer.md) (`/v1/chat/completions/derender`)
|
||||
- Derender chat completion requests
|
||||
- [Completions Derender API](derenderer.md) (`/v1/completions/derender`)
|
||||
- Derender completion requests
|
||||
|
||||
## Tokenize APIs
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# Derenderer APIs
|
||||
|
||||
The derenderer API is the post processing counterpart to the [Renderer APIs](renderer.md). Where `/render` turns a request into token ID (preprocessing), `/derender` turns generated token IDs back into a fully formed OpenAI compatible response (detokenization, reasoning parsing, tool call parsing), all without a GPU.
|
||||
|
||||
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
|
||||
- **Non-streaming**: The endpoints expect a complete `GenerateResponse` with all token IDs present and perform one-shot parsing. Streaming derender would require a separate endpoint design and is not currently supported but is in the pipeline
|
||||
|
||||
Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render`
|
||||
endpoints.
|
||||
|
||||
## Pipeline
|
||||
|
||||
```text
|
||||
render generate derender
|
||||
request ───────────────▶ token_ids ─────────▶ token_ids ──────────▶ response
|
||||
(chat / (GPU less) (token-in / (GPU less) (OpenAI
|
||||
completion) │ token-out engine) ▲ compatible)
|
||||
└─────────────── request + prompt_tokens ──┘
|
||||
```
|
||||
|
||||
The derender step needs more than the engine's `token_ids`. It also consumes the original `chat_request`/`completion_request` and `prompt_tokens` carried over from the render step (see [Request format](#request-format)) so the tool and reasoning parsers have the context they need.
|
||||
|
||||
## API Reference
|
||||
|
||||
- Chat Completions Derender API (`/v1/chat/completions/derender`)
|
||||
- Post process a single `GenerateResponse` into a `ChatCompletionResponse`
|
||||
- Completions Derender API (`/v1/completions/derender`)
|
||||
- Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse`
|
||||
|
||||
## Request format
|
||||
|
||||
Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU.
|
||||
|
||||
`/v1/chat/completions/derender`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-request"
|
||||
```
|
||||
|
||||
`/v1/completions/derender`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request"
|
||||
```
|
||||
|
||||
Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs.
|
||||
|
||||
## Example
|
||||
|
||||
The example below drives the full `render → generate → derender` round trip for a chat request against a GPU less render server (`/render`, `/derender`) and a token-in / token-out engine (`/inference/v1/generate`).
|
||||
|
||||
```python
|
||||
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,
|
||||
}
|
||||
|
||||
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_tokens = len(generate_request["token_ids"])
|
||||
|
||||
# 2. Generate: token IDs -> token IDs (token-in / token-out engine)
|
||||
generate_response = client.post(
|
||||
f"{ENGINE}/inference/v1/generate", json=generate_request
|
||||
).json()
|
||||
|
||||
# 3. Derender: token IDs -> ChatCompletionResponse (GPU less)
|
||||
response = client.post(
|
||||
f"{RENDER}/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": MODEL,
|
||||
"generate_response": generate_response,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"chat_request": chat_request,
|
||||
},
|
||||
).json()
|
||||
|
||||
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. Omit `chat_request` for plain detokenization only.
|
||||
@@ -12,3 +12,5 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and
|
||||
- Render completion requests
|
||||
- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`)
|
||||
- Render chat completions
|
||||
|
||||
For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md).
|
||||
|
||||
@@ -155,8 +155,10 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok
|
||||
- `/v1/chat/completions` - Chat completions
|
||||
- `/v1/chat/completions/batch` - Batch chat completions
|
||||
- `/v1/chat/completions/render` - Render chat completion requests
|
||||
- `/v1/chat/completions/derender` - Derender chat completion requests
|
||||
- `/v1/completions` - Text completions
|
||||
- `/v1/completions/render` - Render completion requests
|
||||
- `/v1/completions/derender` - Derender completion requests
|
||||
- `/v1/embeddings` - Generate embeddings
|
||||
- `/v1/audio/transcriptions` - Audio transcription
|
||||
- `/v1/audio/translations` - Audio translation
|
||||
|
||||
@@ -210,8 +210,39 @@ async def stream_decode_response(session, response, request_id):
|
||||
await session.close()
|
||||
|
||||
|
||||
def example_round_robin_dp_loader(request_number, dp_size):
|
||||
return request_nums % dp_size
|
||||
def flat_interleaved_dp_route(request_number, instances):
|
||||
"""Flat round-robin over the full (instance, dp_rank) slot space.
|
||||
|
||||
ONE counter over (n_instances * dp_size) slots, so instance-selection and
|
||||
DP-rank-selection are derived from the SAME index and can never alias. The
|
||||
previous scheme computed instance = req % n and rank = req % dp from the
|
||||
same counter with n | dp, which locked each instance to a stride-n subset
|
||||
of its ranks (e.g. 2 prefill instances -> 4 of 8 ranks each -> half the
|
||||
GPUs never receive a request, so the deployment falsely appears not to
|
||||
scale).
|
||||
|
||||
Interleaved order — inst0_r0, inst1_r0, inst0_r1, inst1_r1, ... — so
|
||||
consecutive requests alternate instances AND every rank gets walked.
|
||||
|
||||
Assumes homogeneous dp_size across a role's instances (true for the
|
||||
DP<->DP and DP<->TP deployments this proxy targets). Returns
|
||||
(instance_index, dp_rank); dp_rank is None when dp_size == 1 (e.g. a TP
|
||||
decode), which avoids forwarding an out-of-range data-parallel rank.
|
||||
"""
|
||||
n = len(instances)
|
||||
dp = instances[0]["dp_size"]
|
||||
slot = (request_number - 1) % (n * dp)
|
||||
inst_idx = slot % n
|
||||
dp_rank = (slot // n) if dp > 1 else None
|
||||
return inst_idx, dp_rank
|
||||
|
||||
|
||||
@app.route("/health", methods=["GET"])
|
||||
async def health():
|
||||
# Benchmark harnesses and load balancers often probe the proxy URL and
|
||||
# treat a non-200 response as a dead service. Backends expose /health, but
|
||||
# this proxy used to lack one; report the proxy process itself as healthy.
|
||||
return ("ok", 200)
|
||||
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
@@ -244,18 +275,21 @@ async def handle_request(api: str, request: Request):
|
||||
503,
|
||||
)
|
||||
)
|
||||
pid = request_nums % len(prefill_instances)
|
||||
did = request_nums % len(decode_instances)
|
||||
# Flat interleaved round-robin (see flat_interleaved_dp_route): ONE
|
||||
# counter over the full (instance, dp_rank) slot space per role, so
|
||||
# instance-selection and DP-rank-selection derive from the same index
|
||||
# and can never alias. The old scheme keyed both on request_nums with
|
||||
# n_instances | dp_size, stranding half the ranks (e.g. in 2P_DP8EP).
|
||||
pid, selected_prefill_dp_rank = flat_interleaved_dp_route(
|
||||
request_nums, prefill_instances
|
||||
)
|
||||
# Decode instance selection uses the same interleaved walk; in READ
|
||||
# mode the decode reads KV from selected_prefill_dp_rank, so the
|
||||
# decode's own dp_rank is not forwarded here.
|
||||
did, _ = flat_interleaved_dp_route(request_nums, decode_instances)
|
||||
prefill_instance_endpoint = prefill_instances[pid]
|
||||
decode_instance_endpoint = decode_instances[did]
|
||||
|
||||
selected_prefill_dp_rank = None
|
||||
if prefill_instance_endpoint["dp_size"] > 1:
|
||||
selected_prefill_dp_rank = example_round_robin_dp_loader(
|
||||
request_nums // len(prefill_instance_endpoint),
|
||||
prefill_instance_endpoint["dp_size"],
|
||||
)
|
||||
|
||||
# Embed both zmq_addresses in the request_id so the connector can parse
|
||||
# the peer's host/ports from it, similar to P2P-NCCL
|
||||
uid = str(uuid.uuid4()).replace("-", "")
|
||||
@@ -288,7 +322,13 @@ async def handle_request(api: str, request: Request):
|
||||
)
|
||||
)
|
||||
|
||||
req_data["max_tokens"] -= 1
|
||||
# max_completion_tokens takes precedence when present (that's the limit
|
||||
# the backend enforces); fall back to max_tokens. If neither is set
|
||||
# (e.g. benchmark clients that omit it), leave the request unchanged.
|
||||
if "max_completion_tokens" in req_data:
|
||||
req_data["max_completion_tokens"] -= 1
|
||||
elif "max_tokens" in req_data:
|
||||
req_data["max_tokens"] -= 1
|
||||
|
||||
req_data["kv_transfer_params"] = {
|
||||
"do_remote_decode": False,
|
||||
@@ -413,9 +453,33 @@ if __name__ == "__main__":
|
||||
args = parser.parse_args()
|
||||
|
||||
t = start_service_discovery("0.0.0.0", 36367)
|
||||
app.debug = True
|
||||
# High-concurrency hardening. Quart's app.run() uses a shallow listen
|
||||
# backlog (100) and, with app.debug=True, adds per-request overhead that
|
||||
# starves the single accept loop. Under a burst of ~512 simultaneous client
|
||||
# connections the backlog overflows and the kernel RSTs the excess, so
|
||||
# clients see "ClientOSError: [Errno 104] Connection reset by peer" before
|
||||
# any response (~16% request loss at c=512). Serve via hypercorn with debug
|
||||
# OFF and a deep backlog so the burst QUEUES (higher TTFT) instead of being
|
||||
# reset -> 100% request success.
|
||||
app.debug = False
|
||||
app.config["BODY_TIMEOUT"] = 360000
|
||||
app.config["RESPONSE_TIMEOUT"] = 360000
|
||||
|
||||
app.run(host="0.0.0.0", port=args.port)
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from hypercorn.asyncio import serve as _hypercorn_serve
|
||||
from hypercorn.config import Config as _HypercornConfig
|
||||
|
||||
_hcfg = _HypercornConfig()
|
||||
_hcfg.bind = [f"0.0.0.0:{args.port}"]
|
||||
# Deep listen backlog so a wide connection burst queues, not RSTs. NOTE:
|
||||
# effective backlog is capped by the host's net.core.somaxconn (proxy runs
|
||||
# --network host); kernel 6.x defaults to 4096. Override via
|
||||
# PROXY_LISTEN_BACKLOG.
|
||||
_hcfg.backlog = int(os.environ.get("PROXY_LISTEN_BACKLOG", "4096"))
|
||||
# Long-lived SSE streams (8k1k decode ~5 min): never reap on keepalive.
|
||||
_hcfg.keep_alive_timeout = 360000.0
|
||||
|
||||
asyncio.run(_hypercorn_serve(app, _hcfg))
|
||||
t.join()
|
||||
|
||||
@@ -42,12 +42,16 @@ class BlockStored(KVCacheEvent):
|
||||
"""
|
||||
|
||||
group_idx: int | None = None
|
||||
kv_cache_spec_kind: str | None = None
|
||||
kv_cache_spec_sliding_window: int | None = None
|
||||
locality: str | None = None
|
||||
|
||||
|
||||
class BlockRemoved(KVCacheEvent):
|
||||
block_hashes: list[ExternalBlockHash]
|
||||
medium: str | None
|
||||
group_idx: int | None = None
|
||||
locality: str | None = None
|
||||
|
||||
|
||||
class AllBlocksCleared(KVCacheEvent):
|
||||
|
||||
@@ -17,6 +17,7 @@ from transformers import AutoProcessor, AutoTokenizer
|
||||
from vllm import LLM, EngineArgs, SamplingParams
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
QUESTION = "What is the content of each image?"
|
||||
@@ -1443,6 +1444,8 @@ def run_generate(
|
||||
engine_args.seed = seed
|
||||
if tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = tensor_parallel_size
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
@@ -1484,6 +1487,8 @@ def run_chat(
|
||||
engine_args.seed = seed
|
||||
if tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = tensor_parallel_size
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
sampling_params = (
|
||||
|
||||
@@ -21,6 +21,7 @@ from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -2646,6 +2647,8 @@ def main(args):
|
||||
if args.tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = args.tensor_parallel_size
|
||||
engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
# Don't want to check the flag multiple times, so just hijack `prompts`.
|
||||
|
||||
+2
-1
@@ -122,7 +122,8 @@ python = "./.venv"
|
||||
[tool.typos.files]
|
||||
# these files may be written in non english words
|
||||
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
|
||||
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
|
||||
"benchmarks/sonnet.txt", "rust/src/bench/src/datasets/sonnet.txt",
|
||||
"tests/lora/data/*", "build/*",
|
||||
"examples/pooling/token_embed/*", "tests/models/language/pooling/*",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
|
||||
"tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
|
||||
|
||||
@@ -23,8 +23,6 @@ tiktoken >= 0.6.0 # Required for DBRX tokenizer
|
||||
lm-format-enforcer == 0.11.3
|
||||
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x"
|
||||
outlines_core == 0.2.14
|
||||
# required for outlines backend disk cache
|
||||
diskcache == 5.6.3
|
||||
lark == 1.2.2
|
||||
xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
typing_extensions >= 4.10
|
||||
|
||||
@@ -11,9 +11,12 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
|
||||
torchcodec >= 0.14
|
||||
PyNvVideoCodec==2.0.4
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
flashinfer-python==0.6.13
|
||||
flashinfer-cubin==0.6.13
|
||||
apache-tvm-ffi==0.1.9
|
||||
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
|
||||
# install_requires so the published wheel does not carry an unresolvable pin
|
||||
--extra-index-url https://flashinfer.ai/whl/
|
||||
flashinfer-python==0.6.14
|
||||
flashinfer-cubin==0.6.14
|
||||
apache-tvm-ffi==0.1.10
|
||||
tilelang==0.1.9
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
# Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
|
||||
@@ -22,8 +25,8 @@ nvtx==0.2.15
|
||||
fastsafetensors >= 0.3.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.5.2
|
||||
quack-kernels>=0.3.3
|
||||
nvidia-cutlass-dsl[cu13]==4.6.0
|
||||
quack-kernels>=0.4.0 # Required for tml-fa4
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
@@ -2,5 +2,5 @@ lmcache >= 0.3.9
|
||||
# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0
|
||||
# until a fixed newer release is verified for runtime images.
|
||||
cupy-cuda13x < 14.1.0
|
||||
nixl == 1.3.0
|
||||
nixl == 1.3.1
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -21,7 +21,7 @@ conch-triton-kernels==1.2.1
|
||||
timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
# To be consistent with test_quark.py
|
||||
amd-quark>=0.8.99
|
||||
amd-quark==0.12.post1
|
||||
tilelang==0.1.10
|
||||
# Required apache-tvm-ffi matching tilelang version
|
||||
apache-tvm-ffi==0.1.10
|
||||
|
||||
@@ -179,8 +179,6 @@ dill==0.3.8
|
||||
# evaluate
|
||||
# lm-eval
|
||||
# multiprocess
|
||||
diskcache==5.6.3
|
||||
# via -r requirements/test/../common.txt
|
||||
distlib==0.3.9
|
||||
# via virtualenv
|
||||
distro==1.9.0
|
||||
@@ -381,7 +379,7 @@ inflect==5.6.2
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/cuda.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
|
||||
@@ -58,7 +58,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5; platform_machine == "x86_64"
|
||||
instanttensor>=0.1.9; platform_machine == "x86_64"
|
||||
decord==0.6.0; platform_machine == "x86_64"
|
||||
# terratorch is temporarily disabled while PyPI has the `lightning` package
|
||||
# in `quarantined` status (every published terratorch version transitively
|
||||
|
||||
@@ -43,7 +43,7 @@ anyio==4.14.1
|
||||
# sse-starlette
|
||||
# starlette
|
||||
# watchfiles
|
||||
apache-tvm-ffi==0.1.9
|
||||
apache-tvm-ffi==0.1.10
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# xgrammar
|
||||
@@ -194,10 +194,6 @@ dill==0.3.8
|
||||
# evaluate
|
||||
# lm-eval
|
||||
# multiprocess
|
||||
diskcache==5.6.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
distlib==0.3.9
|
||||
# via virtualenv
|
||||
distro==1.9.0
|
||||
@@ -402,7 +398,7 @@ inflect==5.6.2
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/cuda.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
|
||||
@@ -44,5 +44,5 @@ numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
instanttensor>=0.1.9
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
@@ -54,7 +54,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
|
||||
numba==0.65.0 # Required for N-gram speculative decoding
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
instanttensor>=0.1.9
|
||||
decord==0.6.0
|
||||
|
||||
# Prithvi tests
|
||||
|
||||
@@ -187,10 +187,6 @@ dill==0.3.8
|
||||
# evaluate
|
||||
# lm-eval
|
||||
# multiprocess
|
||||
diskcache==5.6.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
distlib==0.4.0
|
||||
# via virtualenv
|
||||
distro==1.9.0
|
||||
@@ -395,7 +391,7 @@ inflect==7.5.0
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.3.0
|
||||
# via pytest
|
||||
instanttensor==0.1.6
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/rocm.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
|
||||
@@ -128,10 +128,6 @@ dill==0.4.1
|
||||
# evaluate
|
||||
# lm-eval
|
||||
# multiprocess
|
||||
diskcache==5.6.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
distro==1.9.0
|
||||
# via
|
||||
# anthropic
|
||||
|
||||
@@ -18,4 +18,4 @@ torchvision
|
||||
torchcodec >= 0.14 # Required for the torchcodec video decoding backend
|
||||
|
||||
auto_round_lib==0.14.1
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# Kill a hung test (per-test 120s budget) so a broken wait fails instead of stalling.
|
||||
[profile.default]
|
||||
slow-timeout = { period = "60s", terminate-after = 2 }
|
||||
Generated
+458
-11
@@ -182,7 +182,7 @@ dependencies = [
|
||||
"hex",
|
||||
"hmac",
|
||||
"rand 0.9.2",
|
||||
"reqwest",
|
||||
"reqwest 0.12.28",
|
||||
"reqwest-eventsource",
|
||||
"secrecy",
|
||||
"serde",
|
||||
@@ -379,6 +379,12 @@ version = "0.22.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
|
||||
|
||||
[[package]]
|
||||
name = "base64ct"
|
||||
version = "1.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
|
||||
|
||||
[[package]]
|
||||
name = "bit-set"
|
||||
version = "0.5.3"
|
||||
@@ -665,6 +671,35 @@ version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b"
|
||||
|
||||
[[package]]
|
||||
name = "cookie"
|
||||
version = "0.18.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
|
||||
dependencies = [
|
||||
"percent-encoding",
|
||||
"time",
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cookie_store"
|
||||
version = "0.22.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206"
|
||||
dependencies = [
|
||||
"cookie",
|
||||
"document-features",
|
||||
"idna",
|
||||
"indexmap 2.13.0",
|
||||
"log",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"time",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.9.4"
|
||||
@@ -888,6 +923,16 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d"
|
||||
dependencies = [
|
||||
"pem-rfc7468",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.8"
|
||||
@@ -1171,6 +1216,12 @@ dependencies = [
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "extended"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "af9673d8203fcb076b19dfd17e38b3d4ae9f44959416ea532ce72415a6020365"
|
||||
|
||||
[[package]]
|
||||
name = "fancy-regex"
|
||||
version = "0.13.0"
|
||||
@@ -1589,13 +1640,15 @@ dependencies = [
|
||||
"indicatif",
|
||||
"libc",
|
||||
"log",
|
||||
"native-tls",
|
||||
"num_cpus",
|
||||
"rand 0.9.2",
|
||||
"reqwest",
|
||||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"ureq",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
@@ -2167,22 +2220,28 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
|
||||
[[package]]
|
||||
name = "llm-multimodal"
|
||||
version = "1.7.1"
|
||||
source = "git+https://github.com/smg-project/llm-multimodal?rev=c8a29dcc755139fdc26185f400ea48c6d6d48273#c8a29dcc755139fdc26185f400ea48c6d6d48273"
|
||||
source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"base64 0.22.1",
|
||||
"blake3",
|
||||
"bytes",
|
||||
"cc",
|
||||
"fast_image_resize",
|
||||
"hf-hub",
|
||||
"image",
|
||||
"libloading",
|
||||
"ndarray 0.17.2",
|
||||
"once_cell",
|
||||
"reqwest",
|
||||
"pkg-config",
|
||||
"rayon",
|
||||
"realfft",
|
||||
"reqwest 0.13.4",
|
||||
"rustfft",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"symphonia",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
@@ -2526,6 +2585,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2740,6 +2800,15 @@ dependencies = [
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pem-rfc7468"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "percent-encoding"
|
||||
version = "2.3.2"
|
||||
@@ -3237,6 +3306,16 @@ dependencies = [
|
||||
"getrandom 0.3.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_distr"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
"rand 0.9.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
@@ -3245,9 +3324,9 @@ checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||
|
||||
[[package]]
|
||||
name = "rayon"
|
||||
version = "1.11.0"
|
||||
version = "1.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f"
|
||||
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
|
||||
dependencies = [
|
||||
"either",
|
||||
"rayon-core",
|
||||
@@ -3346,6 +3425,12 @@ dependencies = [
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex-lite"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
|
||||
|
||||
[[package]]
|
||||
name = "regex-syntax"
|
||||
version = "0.8.10"
|
||||
@@ -3361,6 +3446,7 @@ dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
@@ -3392,7 +3478,43 @@ dependencies = [
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"wasm-streams 0.4.2",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.13.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-tls",
|
||||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"native-tls",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"rustls-pki-types",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams 0.5.0",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
@@ -3408,7 +3530,7 @@ dependencies = [
|
||||
"mime",
|
||||
"nom",
|
||||
"pin-project-lite",
|
||||
"reqwest",
|
||||
"reqwest 0.12.28",
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
@@ -3439,6 +3561,15 @@ dependencies = [
|
||||
"rustc-hash 2.1.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rlimit"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f35ee2729c56bb610f6dba436bf78135f728b7373bdffae2ec815b2d3eb98cc3"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rmp"
|
||||
version = "0.8.15"
|
||||
@@ -3526,7 +3657,9 @@ version = "0.23.37"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4"
|
||||
dependencies = [
|
||||
"log",
|
||||
"once_cell",
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"rustls-webpki",
|
||||
"subtle",
|
||||
@@ -4052,6 +4185,17 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "socks"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"libc",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spm_precompiled"
|
||||
version = "0.1.4"
|
||||
@@ -4127,6 +4271,203 @@ version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
|
||||
|
||||
[[package]]
|
||||
name = "symphonia"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1758d6c853020a7244de03cc3e0185eaea3f58715122422dd3cc7452e6d4c16a"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"symphonia-bundle-flac",
|
||||
"symphonia-bundle-mp3",
|
||||
"symphonia-codec-aac",
|
||||
"symphonia-codec-adpcm",
|
||||
"symphonia-codec-alac",
|
||||
"symphonia-codec-pcm",
|
||||
"symphonia-codec-vorbis",
|
||||
"symphonia-core",
|
||||
"symphonia-format-caf",
|
||||
"symphonia-format-isomp4",
|
||||
"symphonia-format-mkv",
|
||||
"symphonia-format-ogg",
|
||||
"symphonia-format-riff",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-bundle-flac"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee69ad01236a67260b82fd1ff9790dd75ead29f2f46af145e63b7e72273e0e03"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-bundle-mp3"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "350f1f2f2e19ad4dd315db94304d1eb361b29af070681f94e51b8fdaad769546"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"log",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-codec-aac"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1979c515a76371b186aad2feff5f23e21cbec775bf95de08bf1e3af92a2ad76"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-codec-adpcm"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ebbdfd76d6cc5a601c6292a44357c5b7c82f2cd7cdc0f171421f5c5cff0ea1f"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-codec-alac"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a149cbfc7fb5c405d123a273227d31de17138419552112bf1aa7b73e65827b8"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-codec-pcm"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50baee168f0e9dcf6ba7fc06e8b57eb62072a4490cc7cf13af77e72baae5d328"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-codec-vorbis"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "45b07b4423cd8e0fc472575909a5554b12c2f58e3c190b38c24f042e732fd8de"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-common"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8257891ffa7f05e02b58f4761e2abf7e5278c8744fd59e981559e050f86eef55"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-core",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-core"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "95ec293b5f288383b72a7bffcade6b2860b642cf66f28b3bd5967349a49938b1"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bytemuck",
|
||||
"lazy_static",
|
||||
"log",
|
||||
"num-complex",
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-format-caf"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cde3ca76633d3400ab57195456c09f8a58d775ff5452329f3f212b6efc8622f5"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-format-isomp4"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d179a01305b3505940135a9f0180d6ef4b487912748fe97554756f120fbd05e"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-format-mkv"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fb17713e134f5ad316c2690fa3104590ccc85842cdbcf82c3cd1a845cb08aa74"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-format-ogg"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b05a67e02b1e4fca1a261ba4fe06910a9357489ad8c36aafdd2960e9c6559433"
|
||||
dependencies = [
|
||||
"log",
|
||||
"symphonia-common",
|
||||
"symphonia-core",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-format-riff"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "17424452a777666d3eaf09a5c651029b15b6a333812fcc5b5474f2a3f0cff3f0"
|
||||
dependencies = [
|
||||
"extended",
|
||||
"log",
|
||||
"symphonia-core",
|
||||
"symphonia-metadata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "symphonia-metadata"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a31acf5cd623398a6208e2225d18f4b20f761c55098a796a5247ad516a4a8681"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
"log",
|
||||
"regex-lite",
|
||||
"smallvec",
|
||||
"symphonia-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "1.0.109"
|
||||
@@ -4982,6 +5323,42 @@ version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
|
||||
|
||||
[[package]]
|
||||
name = "ureq"
|
||||
version = "3.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"cookie_store",
|
||||
"der",
|
||||
"flate2",
|
||||
"log",
|
||||
"native-tls",
|
||||
"percent-encoding",
|
||||
"rustls",
|
||||
"rustls-pki-types",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"socks",
|
||||
"ureq-proto",
|
||||
"utf8-zero",
|
||||
"webpki-root-certs",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ureq-proto"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"http",
|
||||
"httparse",
|
||||
"log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "url"
|
||||
version = "2.5.8"
|
||||
@@ -5006,6 +5383,12 @@ version = "1.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246"
|
||||
|
||||
[[package]]
|
||||
name = "utf8-zero"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e"
|
||||
|
||||
[[package]]
|
||||
name = "utf8_iter"
|
||||
version = "1.0.4"
|
||||
@@ -5077,6 +5460,38 @@ version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "vllm-bench"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"clap",
|
||||
"dirs",
|
||||
"futures",
|
||||
"hf-hub",
|
||||
"image",
|
||||
"indicatif",
|
||||
"mimalloc",
|
||||
"rand 0.9.2",
|
||||
"rand_distr",
|
||||
"rayon",
|
||||
"reqwest 0.12.28",
|
||||
"rlimit",
|
||||
"rustc-hash 1.1.0",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tiktoken-rs 0.9.1",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-chat"
|
||||
version = "0.1.0"
|
||||
@@ -5097,7 +5512,7 @@ dependencies = [
|
||||
"ndarray 0.17.2",
|
||||
"oss-harmony",
|
||||
"paste",
|
||||
"reqwest",
|
||||
"reqwest 0.13.4",
|
||||
"rmp-serde",
|
||||
"serde",
|
||||
"serde-json-fmt",
|
||||
@@ -5145,6 +5560,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"vllm-bench",
|
||||
"vllm-chat",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-managed-engine",
|
||||
@@ -5344,7 +5760,7 @@ dependencies = [
|
||||
"futures",
|
||||
"hf-hub",
|
||||
"itertools 0.14.0",
|
||||
"reqwest",
|
||||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
@@ -5368,7 +5784,7 @@ dependencies = [
|
||||
"criterion",
|
||||
"fastokens",
|
||||
"hf-hub",
|
||||
"reqwest",
|
||||
"reqwest 0.12.28",
|
||||
"riptoken",
|
||||
"rustc-hash 1.1.0",
|
||||
"serde",
|
||||
@@ -5531,6 +5947,19 @@ dependencies = [
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasmparser"
|
||||
version = "0.244.0"
|
||||
@@ -5563,6 +5992,24 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-root-certs"
|
||||
version = "1.0.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267"
|
||||
dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-roots"
|
||||
version = "1.0.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf"
|
||||
dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "weezl"
|
||||
version = "0.1.12"
|
||||
|
||||
+12
-1
@@ -1,5 +1,6 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"src/bench",
|
||||
"src/chat",
|
||||
"src/cmd",
|
||||
"src/engine-core-client",
|
||||
@@ -32,8 +33,10 @@ base64 = "0.22.1"
|
||||
bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }
|
||||
byteorder = "1.5.0"
|
||||
bytes = "1.12.0"
|
||||
chrono = "0.4.42"
|
||||
clap = { version = "4.5.38", features = ["derive", "env"] }
|
||||
criterion = "0.5.1"
|
||||
dirs = "6.0.0"
|
||||
easy-ext = "1.0.3"
|
||||
educe = "0.6.0"
|
||||
enum-as-inner = "0.7.0"
|
||||
@@ -50,10 +53,12 @@ hyper-util = { version = "0.1.20", features = [
|
||||
"service",
|
||||
"tokio",
|
||||
] }
|
||||
image = { version = "0.25.9", default-features = false, features = ["jpeg"] }
|
||||
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 = "c8a29dcc755139fdc26185f400ea48c6d6d48273" }
|
||||
llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
|
||||
mimalloc = "0.1.52"
|
||||
minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
|
||||
minijinja-contrib = { version = "2.0", features = ["pycompat"] }
|
||||
@@ -71,9 +76,13 @@ prost-types = "0.14.3"
|
||||
pyo3 = "0.28.3"
|
||||
pythonize = "0.28.0"
|
||||
rand = "0.9.2"
|
||||
rand_distr = "0.5.1"
|
||||
rayon = "1.11.0"
|
||||
reasoning-parser = "1.2.2"
|
||||
reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] }
|
||||
reqwest-0-13 = { package = "reqwest", version = "0.13.4", default-features = false, features = ["native-tls"] }
|
||||
riptoken = { version = "0.3.0", default-features = false }
|
||||
rlimit = "0.11.0"
|
||||
rmp-serde = "1.3.1"
|
||||
rmpv = { version = "1.3.1", features = ["with-serde"] }
|
||||
rustc-hash = "1.1.0"
|
||||
@@ -120,8 +129,10 @@ tracing = { version = "0.1.44", features = ["release_max_level_debug"] }
|
||||
tracing-futures = { version = "0.2.5", features = ["futures-03"] }
|
||||
tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] }
|
||||
trait-set = "0.3.0"
|
||||
url = "2.5.7"
|
||||
uuid = { version = "1.22.0", features = ["v4"] }
|
||||
validator = { version = "0.20.0", features = ["derive"] }
|
||||
vllm-bench = { path = "src/bench" }
|
||||
vllm-chat = { path = "src/chat" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
# AGENTS.md
|
||||
|
||||
## Project Overview
|
||||
|
||||
Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime.
|
||||
|
||||
Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend.
|
||||
|
||||
## Build & Test
|
||||
|
||||
Run from the `rust/` workspace root:
|
||||
|
||||
```bash
|
||||
# Build release binary (rust/target/release/vllm-bench)
|
||||
cargo build -p vllm-bench --release
|
||||
|
||||
# Run all tests
|
||||
cargo test -p vllm-bench
|
||||
|
||||
# Run ignored integration tests (requires network for tokenizer download)
|
||||
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/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)
|
||||
- `src/benchmark.rs` — Core benchmark orchestrator (spawn-per-request with tokio + Semaphore; fetches speculative decoding metrics from `/metrics`)
|
||||
- `src/multi_turn.rs` — Multi-turn conversation orchestrator (channel-based worker pool, sequential turns per conversation)
|
||||
- `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/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`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
|
||||
- `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up
|
||||
- `src/ready_checker.rs` — Endpoint readiness with retry
|
||||
- `src/backends/` — Backend implementations (enum dispatch, not trait objects)
|
||||
- `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn)
|
||||
- `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments
|
||||
- `openai_completions.rs` — `/v1/completions` backend
|
||||
- `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal)
|
||||
- `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`
|
||||
- `src/datasets/random.rs` — Random dataset generation with rayon parallelism
|
||||
- `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`)
|
||||
- `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching
|
||||
- `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`)
|
||||
- `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache)
|
||||
- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection)
|
||||
- `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template)
|
||||
- `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`)
|
||||
- `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`)
|
||||
- `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`)
|
||||
- `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs
|
||||
- `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics`
|
||||
- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop)
|
||||
- `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown
|
||||
- `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics`
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
- **Enum dispatch** for backends (avoids async trait object issues with `dyn`)
|
||||
- **reqwest http1_only()** to match Python aiohttp behavior
|
||||
- **rayon** for parallel dataset generation (key perf win over Python)
|
||||
- **mimalloc** global allocator to reduce contention at 1400+ concurrency (page-agnostic; works on aarch64 64K-page kernels where jemalloc aborts with `LG_PAGE=12` builds)
|
||||
- **Arc\<str\> prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs)
|
||||
- **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern)
|
||||
- **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split
|
||||
- **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback.
|
||||
- **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub
|
||||
- **Pre-serialized mm fragments** (`Arc<str>`) for multimodal: image content stored as JSON strings, zero-copy concatenated into payload — avoids deep-cloning ~200KB+ base64 per request
|
||||
- **Steady-state metrics** (default-on in closed-loop): measure throughput/TTFT/TPOT only over the saturated plateau to cut run-to-run variance at high concurrency; `steady_state` is an `Option` in JSON (`#[serde(default)]` for backward compat), null when the scope gate fails or `--no-steady-state`
|
||||
- **`--prompt-token-ids`** (random dataset only): send token-ID arrays instead of text to skip server-side tokenization; also skips the token-length verification pass (counts exact by construction)
|
||||
- **`--random-range-ratio`** follows Python semantics: lengths sampled uniformly from `[len*(1-r), len*(1+r)]`, default `0.0` = fixed; accepts a float in `[0,1)` or `'{"input": r1, "output": r2}'`. (The pre-2026-07 Rust-only form `[len*r, len]` with default 1.0 is rejected with a migration hint.)
|
||||
- **`prompt_list`** (`Arc<[Arc<str>]>` on `SampleRequest`/`RequestFuncInput`): multiple inputs per request for pooling backends — embeddings batches (`--random-batch-size`) send `"input": [...]`, rerank sends `[0]` as query + `[1..]` as documents
|
||||
- JSON output schema must match Python `vllm bench serve` exactly
|
||||
|
||||
## Common Issues
|
||||
|
||||
- **localhost vs 127.0.0.1**: Some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. Use `127.0.0.1` or the actual hostname.
|
||||
- **Models without tokenizer.json** (e.g., `nvidia/Kimi-K2.5-NVFP4`): Automatically falls back to server-side tokenization. Can also use `--tokenizer` to point to a model with `tokenizer.json`.
|
||||
- **usage.completion_tokens parsing**: vLLM sends final usage chunk with `"choices":[]` (empty array). The usage `if` must be separate from the choices `if` (not `else if`).
|
||||
|
||||
## Typical Usage
|
||||
|
||||
```bash
|
||||
# Embedding benchmark (openai-embeddings, 8 inputs batched per request)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-embeddings \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model BAAI/bge-large-en-v1.5 \
|
||||
--dataset-name random \
|
||||
--random-input-len 512 \
|
||||
--random-batch-size 8 \
|
||||
--num-prompts 1000 \
|
||||
--save-result
|
||||
|
||||
# vLLM rerank benchmark (one query + 8 documents per request)
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm-rerank \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model BAAI/bge-reranker-v2-m3 \
|
||||
--dataset-name random-rerank \
|
||||
--random-input-len 512 \
|
||||
--random-batch-size 8 \
|
||||
--num-prompts 500 \
|
||||
--save-result
|
||||
|
||||
# Prefix-cache stress (10 shared prefixes, 256+256 tokens)
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name prefix_repetition \
|
||||
--prefix-repetition-prefix-len 256 \
|
||||
--prefix-repetition-suffix-len 256 \
|
||||
--prefix-repetition-num-prefixes 10 \
|
||||
--num-prompts 1000
|
||||
|
||||
# Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name custom \
|
||||
--dataset-path workload.jsonl \
|
||||
--custom-output-len -1 \
|
||||
--num-prompts 1000
|
||||
|
||||
# Random dataset
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name random \
|
||||
--random-input-len 8192 \
|
||||
--random-output-len 1024 \
|
||||
--ignore-eos \
|
||||
--num-prompts 4096 \
|
||||
--percentile-metrics "ttft,tpot,itl,e2el" \
|
||||
--save-result \
|
||||
--max-concurrency 1400
|
||||
|
||||
# Random multimodal dataset (VLM benchmark)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--dataset-name random-mm \
|
||||
--random-input-len 512 \
|
||||
--random-output-len 128 \
|
||||
--num-prompts 100 \
|
||||
--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}'
|
||||
|
||||
# HuggingFace dataset (WildChat)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name hf \
|
||||
--dataset-path allenai/WildChat-4.8M \
|
||||
--hf-split train \
|
||||
--num-prompts 1000 \
|
||||
--save-result
|
||||
|
||||
# HuggingFace dataset (LongBench with subset)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name hf \
|
||||
--dataset-path THUDM/LongBench \
|
||||
--hf-subset narrativeqa \
|
||||
--hf-split test \
|
||||
--hf-output-len 512 \
|
||||
--num-prompts 200
|
||||
```
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user