iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into sync/lesj/qwen4exp-bugfixes

Browse Source 
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
This commit is contained in:
lesj0610
2026-09-16 13:41:12 +09:00
parent a66a3e1b4e af1c01499b
commit 21d2dd56d8
517 changed files with 21831 additions and 6611 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/hardware_tests/cpu.yaml
+46
View File
@@ -52,6 +52,52 @@ steps:
tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
parallelism: 2
# Speculative decoding under a CPU platform. The unit tests already pass on
# CPU but no CPU job runs them; the e2e tests gate the engine path itself,
# which unit tests cannot reach.
- label: CPU-Spec Decode Tests
key: cpu-spec-decode-tests
depends_on: []
soft_fail: true
device: intel_cpu
no_plugin: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/cpu/
- vllm/v1/worker/cpu_model_runner.py
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/attention/backends/cpu_attn.py
- vllm/config/speculative.py
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/qwen3_5_mtp.py
- vllm/model_executor/models/llama_eagle3.py
- vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
- vllm/model_executor/layers/mamba/mamba_mixer2.py
- vllm/model_executor/layers/mamba/ops/cpu/
- tests/v1/e2e/test_cpu_spec_decode.py
- tests/v1/spec_decode/
- tests/v1/worker/
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
pytest -x -v -s \
tests/v1/worker/test_gpu_autoregressive_speculator.py \
tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \
tests/v1/worker/test_gpu_batch_ordering.py \
tests/v1/worker/test_gpu_rejection_sampler_chunking.py \
tests/v1/worker/test_eagle3_aux_hidden_states_pp.py \
tests/v1/worker/test_spec_decode_embed_sharing_pp.py \
tests/v1/spec_decode/test_adaptive_verification.py \
tests/v1/spec_decode/test_dflash2.py \
tests/v1/spec_decode/test_dflash_causality.py \
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py \
tests/v1/spec_decode/test_draft_attention_backend_override.py \
tests/v1/spec_decode/test_draft_moe_backend_override.py \
tests/v1/spec_decode/test_ngram.py \
tests/v1/spec_decode/test_llm_base_proposer_sampling.py
unset VLLM_CPU_CI_ENV
pytest -x -v -s tests/v1/e2e/test_cpu_spec_decode.py"
# Note: SDE can't be downloaded from CI host because of AWS WAF
# - label: CPU-Compatibility Tests
# depends_on: []
.buildkite/intel_jobs/entrypoints_intel.yaml
+2 -1
View File
@@ -134,7 +134,8 @@ steps:
'cd tests &&
pytest -v -s tool_use
--deselect="tests/tool_use/mistral/test_mistral_tool_calls.py::test_tool_call_none_with_tools[ministral-3b]"
--deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" &&
--deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]"
--deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[llama3.2]" &&
pytest -v -s entrypoints/tool_parsers &&
pytest -v -s entrypoints/generate &&
pytest -v -s entrypoints/anthropic &&
.buildkite/intel_jobs/lm_eval_intel.yaml
+2 -1
View File
@@ -23,5 +23,6 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
'uv pip install "gpt-oss[eval]==0.0.5" &&
cd tests &&
pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt'
.buildkite/intel_jobs/test-intel.yaml
+1 -1
View File
@@ -97,7 +97,7 @@ steps:
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
'cd tests && pip install dpctl &&
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
.buildkite/scripts/ci-otel/ci_otel.sh
+58 -7
View File
@@ -26,7 +26,8 @@ if [ -z "${CI_INFRA_OTEL_RUNTIME_DIR}" ] ||
[ -z "${_CI_INFRA_OTEL_PYTHON}" ] ||
[ ! -f "${CI_INFRA_OTEL_DIR}/ci_otel.py" ] ||
! mkdir -p "${CI_INFRA_OTEL_SPOOL_DIR}" ||
! PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \
! PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \
"${_CI_INFRA_OTEL_PYTHON}" -c "import ci_otel" >/dev/null 2>&1; then
echo "vLLM CI OTel: tracing unavailable; test command will run normally" >&2 || :
if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
@@ -58,12 +59,45 @@ fi
_ci_otel_python() {
if command -v timeout >/dev/null 2>&1; then
timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@"
PYTHONDONTWRITEBYTECODE=1 timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@"
else
"${_CI_INFRA_OTEL_PYTHON}" "$@"
PYTHONDONTWRITEBYTECODE=1 "${_CI_INFRA_OTEL_PYTHON}" "$@"
fi
}
# The Docker test commands run as root, while /workdir is owned by the
# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so
# any future project-code import cannot leave root-owned files in the checkout.
_ci_otel_prepare_gpu_user() {
_CI_INFRA_GPU_UID=""
_CI_INFRA_GPU_GID=""
[ "$(id -u 2>/dev/null || :)" = "0" ] || return 0
command -v setpriv >/dev/null 2>&1 || return 1
_CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}"
[ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1
# The owner fields are numeric and whitespace-free on GNU stat.
# shellcheck disable=SC2046
set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :)
[ "$#" -eq 2 ] || return 1
case "$1:$2" in
*[!0-9:]* | 0:* | *:0) return 1 ;;
esac
_CI_INFRA_GPU_UID="$1"
_CI_INFRA_GPU_GID="$2"
chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \
"${CI_INFRA_OTEL_SPOOL_DIR}" || return 1
chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1
if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
# Permit traversal to the private child spool, but not directory listing.
chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1
fi
setpriv --reuid="${_CI_INFRA_GPU_UID}" \
--regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \
-- test -w "${CI_INFRA_OTEL_SPOOL_DIR}"
}
ci_otel_start() {
_CI_INFRA_OTEL_COMMAND_INDEX="$1"
_CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}"
@@ -86,8 +120,23 @@ ci_otel_start() {
if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] &&
[ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] &&
command -v nvidia-smi >/dev/null 2>&1; then
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" </dev/null &
_CI_INFRA_GPU_PID=$!
if _ci_otel_prepare_gpu_user; then
if [ -n "${_CI_INFRA_GPU_UID}" ]; then
PYTHONDONTWRITEBYTECODE=1 setpriv \
--reuid="${_CI_INFRA_GPU_UID}" --regid="${_CI_INFRA_GPU_GID}" \
--clear-groups --no-new-privs -- \
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" \
</dev/null &
else
PYTHONDONTWRITEBYTECODE=1 \
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" \
</dev/null &
fi
_CI_INFRA_GPU_PID=$!
else
echo "vLLM CI OTel: non-root GPU sampler unavailable; sampling disabled" \
>&2 || :
fi
fi
}
@@ -137,9 +186,11 @@ _ci_otel_on_exit() {
trap - 0
ci_otel_finish "${_CI_INFRA_OTEL_EXIT_STATUS}" || :
if command -v timeout >/dev/null 2>&1; then
timeout 4s "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
PYTHONDONTWRITEBYTECODE=1 timeout 4s \
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
else
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
PYTHONDONTWRITEBYTECODE=1 \
"${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || :
fi
if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then
rm -rf -- "${CI_INFRA_OTEL_RUNTIME_DIR}" || :
.buildkite/scripts/ci-otel/tests/test_ci_otel.py
+109
View File
@@ -6,6 +6,7 @@ from __future__ import annotations
import json
import os
import shlex
import shutil
import subprocess
import sys
import threading
@@ -396,6 +397,35 @@ def test_shell_wrapper_preserves_failure_status(tmp_path):
assert result.returncode == 1
def test_shell_wrapper_does_not_write_bytecode_to_helper_tree(tmp_path):
helper_dir = tmp_path / "ci-otel"
shutil.copytree(
SCRIPTS_DIR,
helper_dir,
ignore=shutil.ignore_patterns("__pycache__"),
)
shell = (
f'. "{helper_dir / "ci_otel.sh"}"; '
f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0"
)
result = subprocess.run(
["/bin/sh", "-c", shell],
check=False,
capture_output=True,
text=True,
env={
**os.environ,
"CI_INFRA_OTEL_DIR": str(helper_dir),
"CI_INFRA_GPU_SAMPLING": "0",
"CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"),
},
)
assert result.returncode == 0, result.stderr
assert list(helper_dir.rglob("*.pyc")) == []
def test_ci_otel_run_records_command_and_preserves_status(tmp_path):
shell = (
f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; '
@@ -957,6 +987,85 @@ def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch):
assert [len(batch[0].events) for batch in batches] == [1, 1, 1]
def _write_executable(path: Path, body: str) -> None:
path.write_text(f"#!/bin/sh\n{body}")
path.chmod(0o755)
@pytest.mark.parametrize(
"workspace_owner,expect_sampler",
[("997 996", True), ("0 0", False)],
)
def test_root_shell_runs_gpu_sampler_as_checkout_owner(
tmp_path, workspace_owner, expect_sampler
):
bin_dir = tmp_path / "bin"
workspace = tmp_path / "workdir"
runtime = tmp_path / "runtime"
log = tmp_path / "calls.log"
bin_dir.mkdir()
workspace.mkdir()
_write_executable(
bin_dir / "id",
'[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n',
)
_write_executable(
bin_dir / "stat",
f"printf '%s\\n' '{workspace_owner}'\n",
)
for command in ("chown", "chmod"):
_write_executable(bin_dir / command, "exit 0\n")
_write_executable(
bin_dir / "setpriv",
'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n'
'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n'
'[ "$#" -gt 0 ] && shift\n'
'exec "$@"\n',
)
_write_executable(
bin_dir / "python3",
'[ "$1" = "-c" ] && exit 0\n'
'[ "$2" = "new-context" ] && {\n'
' echo "01010101010101010101010101010101 0202020202020202 - 1"\n'
" exit 0\n"
"}\n"
'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n'
"exit 0\n",
)
_write_executable(bin_dir / "nvidia-smi", "exit 0\n")
shell = (
f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; '
f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0"
)
result = subprocess.run(
["/bin/sh", "-c", shell],
check=False,
capture_output=True,
text=True,
env={
**os.environ,
"PATH": f"{bin_dir}:{os.environ['PATH']}",
"CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR),
"CI_INFRA_OTEL_RUNTIME_DIR": str(runtime),
"CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace),
"CI_OTEL_TEST_LOG": str(log),
},
)
assert result.returncode == 0, result.stderr
calls = log.read_text() if log.exists() else ""
if expect_sampler:
expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --"
assert calls.count("setpriv") == 2
assert calls.count(expected) == 2
assert "gpu" in calls
else:
assert calls == ""
assert "non-root GPU sampler unavailable" in result.stderr
@pytest.mark.parametrize("sampling", ["1", "0"])
def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
tmp_path, sampling
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+3 -1
View File
@@ -112,7 +112,9 @@ print_packages
export CONTAINER_NAME
export -f kernel_tests model_tests serving_tests
case "$SHARD_ID" in
0) timeout 30m bash -c kernel_tests ;;
# The kernel shard contains several large parametrized suites and can exceed
# 30 minutes on otherwise healthy Arm runners.
0) timeout 60m bash -c kernel_tests ;;
1) timeout 30m bash -c model_tests ;;
2) timeout 30m bash -c serving_tests ;;
esac
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+6 -3
View File
@@ -59,7 +59,7 @@ docker pull ubuntu:25.04 || true
# building the docker image
echo "--- :docker: Building Docker image"
BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:'
BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output'
BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries
BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3
build_log="$(mktemp)"
@@ -83,8 +83,11 @@ rm -f "$build_log"
# Run the image, setting --shm-size=4g for tensor parallel. Default to
# HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network;
# retry once online if the cache is missing something.
OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)'
# retry once online if the cache is missing something. vllm's get_config()
# wraps the raw huggingface_hub offline-mode errors in a generic ValueError
# (see transformers_utils/config.py), so match that message too or the
# fallback never triggers for a config-lookup cache miss.
OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified'
run_test() {
local hf_offline=$1
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \
.buildkite/scripts/install-kv-offload.sh
+18
View File
@@ -0,0 +1,18 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
set -euo pipefail
if [ -z "${CUDA_VERSION:-}" ]; then
exit 0
fi
# TODO: Use the PyPI wheel once it is published.
uv pip install --system \
"kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0"
# Keep only the NIXL wheel matching the CI image's CUDA runtime.
NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p')
uv pip uninstall --system nixl-cu12 nixl-cu13
uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}"
.buildkite/test-amd.yaml
+6 -6
View File
@@ -1067,7 +1067,7 @@ steps:
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- pytest -v -s entrypoints/scale_out
- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD
- label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" # TBD
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
@@ -1080,9 +1080,10 @@ steps:
- tests/entrypoints/openai
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD
- label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
@@ -1096,7 +1097,6 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" # TBD
timeout_in_minutes: 65
@@ -3217,7 +3217,7 @@ steps:
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- pytest -v -s entrypoints/scale_out
- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD
- label: ":amd: (MI355) Entrypoints Integration (OpenAI API completion)" # TBD
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
@@ -3230,9 +3230,10 @@ steps:
- tests/entrypoints/openai
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD
- label: ":amd: (MI355) Entrypoints Integration (OpenAI API chat_completion)" # TBD
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
@@ -3246,7 +3247,6 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD
timeout_in_minutes: 40
.buildkite/test_areas/disaggregated.yaml
+1
View File
@@ -287,6 +287,7 @@ steps:
PREFILL_BLOCK_SIZE: "256"
DECODE_BLOCK_SIZE: "256"
MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
VLLM_ENGINE_READY_TIMEOUT_S: "900"
VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
# The V4 indexer only supports block-outermost layouts; let it pick.
PREFILLER_KV_LAYOUT: ""
.buildkite/test_areas/distributed.yaml
+1 -1
View File
@@ -208,7 +208,7 @@ steps:
- tests/distributed/test_utils
- tests/distributed/test_engram_dp_shard.py
- vllm/config/engram.py
- vllm/models/deepseek_v4_1/
- vllm/models/deepseek_v41/
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
.buildkite/test_areas/entrypoints.yaml
+7 -7
View File
@@ -73,9 +73,9 @@ steps:
depends_on:
- image-build-amd
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 1)"
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API completion)"
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
key: entrypoints-integration-api-server-openai-completion
timeout_in_minutes: 68
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -84,19 +84,20 @@ steps:
- tests/entrypoints/openai
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
mirror:
amd:
label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)"
label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)"
dind: false
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 2)"
- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API chat_completion)"
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
key: entrypoints-integration-api-server-openai-chat_completion
timeout_in_minutes: 83
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -106,10 +107,9 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
mirror:
amd:
label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)"
label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)"
dind: false
device: mi300_1
timeout_in_minutes: 70
.buildkite/test_areas/lm_eval.yaml
+34
View File
@@ -30,6 +30,22 @@ steps:
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking"
device: h200_35gb
key: lm-eval-watermarking
# Recorded runs took 80–140 seconds
timeout_in_minutes: 10
source_file_dependencies:
- .buildkite/test_areas/lm_eval.yaml
- tests/evals/gsm8k/
- vllm/config/watermarking.py
- vllm/sampling_params.py
- vllm/v1/watermarking/
- vllm/v1/worker/gpu/
autorun_on_main: true
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt
# - label: LM Eval Large Models (4xA100)
# key: lm-eval-large-models-4xa100
# device: a100
@@ -100,6 +116,24 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
- label: ":nvidia: (B200) LM Eval HiSparse Nightly"
key: lm-eval-hisparse-nightly-4xb200
timeout_in_minutes: 60
device: b200-k8s
num_devices: 4
optional: true
source_file_dependencies:
- .buildkite/test_areas/lm_eval.yaml
- tests/evals/gsm8k/
- vllm/v1/hisparse/
- vllm/distributed/kv_transfer/kv_connector/v1/hisparse/
- vllm/v1/attention/backends/mla/
- vllm/model_executor/layers/attention/sparse_mla_attention.py
- vllm/v1/worker/gpu/
- csrc/libtorch_stable/hisparse_kernels.cu
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt
- label: ":nvidia: (B200) LM Eval PCP"
key: lm-eval-pcp-4xb200
timeout_in_minutes: 360
.buildkite/test_areas/misc.yaml
+6 -4
View File
@@ -130,6 +130,7 @@ steps:
source_file_dependencies: *v1-core-kv-metrics-deps
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
mirror:
@@ -141,19 +142,20 @@ steps:
depends_on:
- image-build-amd
- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors"
- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors Shard %N"
device: h200_35gb
key: v1-kv-connectors
timeout_in_minutes: 45
parallelism: 4
source_file_dependencies: *v1-core-kv-metrics-deps
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
mirror:
amd:
label: ":amd: (MI300) V1 KV Connectors"
label: ":amd: (MI300) V1 KV Connectors Shard %N"
dind: false
device: mi300_1
timeout_in_minutes: 45
.github/CODEOWNERS
+8 -8
View File
@@ -34,9 +34,9 @@
# Entrypoints
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337
/vllm/entrypoints/mcp @heheda12345
/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337
/vllm/entrypoints/speech_to_text/realtime @njhill
/vllm/entrypoints/speech_to_text @NickLucche
/vllm/entrypoints/pooling @noooop
@@ -73,7 +73,7 @@
/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
/vllm/v1/sample @22quinn @houseroad @njhill
/vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett @arpera
/vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera
/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium
/vllm/v1/kv_offload @ApostaC @orozery
/vllm/v1/simple_kv_offload @ivanium
@@ -107,8 +107,8 @@
/tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan
/tests/multimodal/media @Isotr0py
/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
/tests/v1/structured_output @mgoin @russellb @aarnphm @arpera
/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera
/tests/v1/structured_output @mgoin @aarnphm @arpera
/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium
/tests/weight_loading @mgoin @youkaichao @yewentao256
/tests/lora @jeejeelee
@@ -232,6 +232,6 @@ mkdocs.yaml @hmellor
/docs/training/weight_transfer @aoshen02
# Security guide and policies
/docs/usage/security.md @russellb @jperezdealgaba
/SECURITY.md @russellb @jperezdealgaba
/docs/contributing/vulnerability_management.md @russellb @jperezdealgaba
/docs/usage/security.md @jperezdealgaba
/SECURITY.md @jperezdealgaba
/docs/contributing/vulnerability_management.md @jperezdealgaba
.github/mergify.yml
+2 -3
View File
@@ -343,9 +343,8 @@ pull_request_rules:
conditions:
- label != stale
- or:
# The 4.1 sources spell the model three ways: deepseek_v4_1 (models),
# deepseek_v41 (most) and deepseekv41 (tool parsers).
- files~=^vllm/models/deepseek_v4_1/
# The 4.1 sources use deepseek_v41 and deepseekv41 (tool parsers).
- files~=^vllm/models/deepseek_v41/
- files~=^vllm/(?:parser|tokenizers|reasoning|transformers_utils/configs)/deepseek_v41.*\.py
- files~=^vllm/tool_parsers/deepseekv41.*\.py
- files~=^rust/src/(?:chat|parser)/src/.*deepseek_v41
.github/workflows/buf.yml
-29
View File
@@ -5,12 +5,6 @@ on:
paths:
- "rust/proto/**"
- ".github/workflows/buf.yml"
push:
tags:
- "v*"
schedule:
- cron: "0 8 * * *"
workflow_dispatch:
permissions:
contents: read
@@ -21,37 +15,14 @@ concurrency:
jobs:
proto:
if: github.repository_owner == 'vllm-project'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
fetch-depth: 0
- uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0
with:
version: 1.72.0
input: rust/proto
token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }}
format: false
breaking: false
pr_comment: false
push: false
- name: Publish nightly schema
if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
run: >-
buf push rust/proto
--create
--create-default-label main
--create-visibility public
--label nightly
--source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}"
- name: Publish release schema
if: github.repository == 'vllm-project/vllm' && github.event_name == 'push'
run: >-
buf push rust/proto
--create
--create-default-label main
--create-visibility public
--label main
--label "${GITHUB_REF_NAME}"
--source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}"
.github/workflows/run-ci-command.yml
+6
View File
@@ -21,11 +21,17 @@ jobs:
(github.event.comment.body == '/ci run' ||
github.event.comment.body == '/ci run all' ||
github.event.comment.body == '/ci run nightly' ||
github.event.comment.body == '/ci run --allow-stale' ||
github.event.comment.body == '/ci run all --allow-stale' ||
github.event.comment.body == '/ci run nightly --allow-stale' ||
github.event.comment.body == '/ci retry' ||
github.event.comment.body == '/ci cancel' ||
github.event.comment.body == '/amd-ci run' ||
github.event.comment.body == '/amd-ci run all' ||
github.event.comment.body == '/amd-ci run nightly' ||
github.event.comment.body == '/amd-ci run --allow-stale' ||
github.event.comment.body == '/amd-ci run all --allow-stale' ||
github.event.comment.body == '/amd-ci run nightly --allow-stale' ||
github.event.comment.body == '/amd-ci retry' ||
github.event.comment.body == '/amd-ci cancel')
runs-on: [self-hosted, linux, x64, vllm-runners]
.github/workflows/scripts/run_ci_command.py
+92 -12
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import http.client
import json
import os
import random
@@ -22,6 +23,7 @@ COMMAND_RUN_AMD_CI_ALL = "/amd-ci run all"
COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly"
COMMAND_RETRY_AMD_FAILED = "/amd-ci retry"
COMMAND_CANCEL_AMD_CI = "/amd-ci cancel"
ALLOW_STALE_SUFFIX = " --allow-stale"
RUN_CI_COMMAND_ENV = {
COMMAND_RUN_CI: {},
COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"},
@@ -30,20 +32,25 @@ RUN_CI_COMMAND_ENV = {
COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"},
COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"},
}
ALLOW_STALE_COMMANDS = frozenset(
f"{command}{ALLOW_STALE_SUFFIX}" for command in RUN_CI_COMMAND_ENV
)
RUN_CI_COMMAND_ENV.update(
{
command: RUN_CI_COMMAND_ENV[command.removesuffix(ALLOW_STALE_SUFFIX)]
for command in ALLOW_STALE_COMMANDS
}
)
UPSTREAM_CI_COMMANDS = frozenset(
{
COMMAND_RUN_CI,
COMMAND_RUN_CI_ALL,
COMMAND_RUN_CI_NIGHTLY,
*(command for command in RUN_CI_COMMAND_ENV if command.startswith("/ci ")),
COMMAND_RETRY_FAILED,
COMMAND_CANCEL_CI,
}
)
AMD_CI_COMMANDS = frozenset(
{
COMMAND_RUN_AMD_CI,
COMMAND_RUN_AMD_CI_ALL,
COMMAND_RUN_AMD_CI_NIGHTLY,
*(command for command in RUN_CI_COMMAND_ENV if command.startswith("/amd-ci ")),
COMMAND_RETRY_AMD_FAILED,
COMMAND_CANCEL_AMD_CI,
}
@@ -79,6 +86,10 @@ class ApiError(RuntimeError):
self.status = status
class CiPreparationError(RuntimeError):
pass
def rate_limit_jitter() -> float:
return random.uniform(1, 5)
@@ -137,6 +148,8 @@ class HttpTransport:
None,
f"API request failed: {error.reason}",
) from error
except (OSError, http.client.HTTPException) as error:
raise ApiError(None, f"API request failed: {error}") from error
if not response_body:
return None
@@ -234,6 +247,17 @@ class GitHubClient:
def get_pr(self, number: int) -> dict[str, Any]:
return self._request(self._repo_path(f"/pulls/{number}"))
def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int:
base_ref = urllib.parse.quote(f"refs/heads/{base_ref}", safe="")
head_sha = urllib.parse.quote(head_sha, safe="")
response = self._request(
self._repo_path(f"/compare/{base_ref}...{head_sha}?per_page=1")
)
behind = response.get("behind_by") if isinstance(response, dict) else None
if type(behind) is not int or behind < 0:
raise ApiError(None, "GitHub returned an invalid commits-behind count.")
return behind
def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]:
commit = urllib.parse.quote(commit, safe="")
response = self._request(self._repo_path(f"/commits/{commit}/pulls"))
@@ -747,6 +771,10 @@ def notify_authorized(
(
f"✅ @{author}, CI is now available for this PR.\n\n"
"- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n"
"- Your branch must contain every commit currently on its upstream "
"target branch. Merge or rebase onto the latest target branch, then "
"rerun the command. Append `--allow-stale` to a run command to test "
"an outdated branch at your own risk.\n"
"- `/ci retry` retries failed jobs in the CI build for the current "
"PR head. If the current head has no CI build, it starts a new CI "
"build for the current head containing only jobs that failed in "
@@ -794,6 +822,56 @@ def resolve_workflow_run_pr(
return find_matching_pr(github.list_pulls_for_commit(head_sha))
def prepare_pr_for_ci(
github: GitHubClient,
pr: Mapping[str, Any],
command: str,
) -> tuple[dict[str, Any], str]:
base_ref = pr["base"]["ref"]
no_action = "No new CI build was started."
try:
behind = github.get_commits_behind_base(base_ref, pr["head"]["sha"])
current_pr = github.get_pr(pr["number"])
except ApiError as error:
raise CiPreparationError(
f"Could not check the PR against upstream `{base_ref}`. "
f"{no_action} Comment `{command}` again. {error}"
) from error
if (
current_pr["state"] != "open"
or current_pr["head"]["sha"] != pr["head"]["sha"]
or current_pr["base"]["ref"] != pr["base"]["ref"]
):
raise CiPreparationError(
f"The PR changed while checking its branch. {no_action} "
f"Comment `{command}` again."
)
commits = "commit" if behind == 1 else "commits"
lag = f"This PR is {behind} {commits} behind upstream `{base_ref}`."
if command in ALLOW_STALE_COMMANDS:
warning = ""
if behind:
run_command = command.removesuffix(ALLOW_STALE_SUFFIX)
warning = (
f"⚠️ {lag} Running CI at your own risk because "
"`--allow-stale` was requested; outdated CI configuration may "
"cause failures. Before merging, merge or rebase onto the latest "
f"`{base_ref}`, then rerun `{run_command}` on the latest PR commit."
)
print(warning)
return current_pr, warning
if behind:
raise CiPreparationError(
f"{lag} Your branch must contain every commit currently on upstream "
f"`{base_ref}`. {no_action} "
f"Merge or rebase onto the latest `{base_ref}`, then rerun `{command}`."
" To test this branch at your own risk, "
f"use `{command}{ALLOW_STALE_SUFFIX}`."
)
return current_pr, ""
def handle_run_ci(
*,
actor: str,
@@ -828,12 +906,7 @@ def handle_run_ci(
f"{ci_name} is already running for this commit: {active_build['web_url']}"
)
current_pr = github.get_pr(pr["number"])
if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
return (
"The PR head changed while processing the command. "
f"Comment `{command}` again."
)
current_pr, warning = prepare_pr_for_ci(github, pr, command)
build = buildkite.create_build(
create_build_payload(
@@ -846,6 +919,7 @@ def handle_run_ci(
return (
f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) "
f"for commit `{current_pr['head']['sha'][:12]}`."
+ (f"\n\n{warning}" if warning else "")
)
@@ -1099,6 +1173,12 @@ def run(
raise ValueError(f"Unsupported CI command: {command}")
add_reaction_safely(github, comment_id, "rocket")
github.add_comment(issue_number, f"✅ {message}")
except CiPreparationError as error:
add_reaction_safely(github, comment_id, "-1")
github.add_comment(
issue_number,
f"❌ {error}\n\n{command_comment_marker(comment_id)}",
)
except Exception:
add_reaction_safely(github, comment_id, "confused")
raise
.github/workflows/scripts/test_run_ci_command.py
+196
View File
@@ -22,8 +22,10 @@ from run_ci_command import (
COMMAND_RUN_CI_ALL,
COMMAND_RUN_CI_NIGHTLY,
RETRY_STATES,
RUN_CI_COMMAND_ENV,
ApiError,
BuildkiteClient,
GitHubClient,
HttpTransport,
authorize,
create_build_payload,
@@ -85,6 +87,7 @@ class FakeGitHub:
comments: list[str] | None = None,
pulls_for_commit: list[dict[str, Any]] | None = None,
prs: dict[int, dict[str, Any]] | None = None,
behind: int = 0,
) -> None:
self.comments = comments or []
self.permission = permission
@@ -95,6 +98,12 @@ class FakeGitHub:
self.reviews = reviews or []
self.pulls_for_commit = pulls_for_commit or []
self.prs = prs or {self.pr["number"]: self.pr}
self.behind = behind
self.lag_queries: list[tuple[str, str]] = []
def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int:
self.lag_queries.append((base_ref, head_sha))
return self.behind
def get_pr(self, number: int) -> dict[str, Any]:
try:
@@ -322,6 +331,11 @@ class RunCiCommandTest(unittest.TestCase):
for command in commands:
with self.subTest(command=command):
self.assertEqual(parse_command(command), command)
override = f"{command} --allow-stale"
if " run" in command:
self.assertEqual(parse_command(override), override)
else:
self.assertIsNone(parse_command(override))
self.assertIsNone(parse_command("/ci run please"))
self.assertIsNone(parse_command("/ci run all please"))
@@ -331,6 +345,8 @@ class RunCiCommandTest(unittest.TestCase):
self.assertIsNone(parse_command("/amd-ci retry "))
self.assertIsNone(parse_command("/AMD-CI run"))
self.assertIsNone(parse_command("/amdci run"))
self.assertIsNone(parse_command("/ci run --allow-stale all"))
self.assertIsNone(parse_command("/amd-ci run --allow-stale please"))
def test_commands_select_only_their_configured_pipeline(self) -> None:
cases = (
@@ -530,6 +546,177 @@ class RunCiCommandTest(unittest.TestCase):
self.assertTrue(github.comments[0].startswith("✅ "))
self.assertIn("Buildkite CI #123", github.comments[0])
def test_run_variants_require_current_base_without_changing_commit(self) -> None:
for command in RUN_CI_COMMAND_ENV:
if command.endswith(" --allow-stale"):
continue
for behind in (0, 1, 5):
with self.subTest(command=command, behind=behind):
github = FakeGitHub(
behind=behind, pr=make_pr(base={"ref": "release/1.0"})
)
buildkite = FakeBuildkite([[], []])
run(make_event(command), github, buildkite)
self.assertEqual(
github.lag_queries,
[("release/1.0", github.pr["head"]["sha"])],
)
if behind == 0:
self.assertEqual(
buildkite.created_builds[0]["commit"],
github.pr["head"]["sha"],
)
self.assertIn(
"](https://buildkite.example/builds/123)",
github.comments[0],
)
else:
self.assertEqual(buildkite.created_builds, [])
self.assertTrue(github.comments[0].startswith("❌ "))
self.assertIn(str(behind), github.comments[0])
self.assertIn("upstream `release/1.0`", github.comments[0])
self.assertIn("every commit", github.comments[0])
self.assertIn(
"No new CI build was started.", github.comments[0]
)
self.assertIn(command, github.comments[0])
def test_allow_stale_preserves_author_permissions_and_run_configuration(
self,
) -> None:
cases = (
(COMMAND_RUN_CI, "ci", {}),
(COMMAND_RUN_CI_ALL, "ci", {"RUN_ALL": "1"}),
(COMMAND_RUN_CI_NIGHTLY, "ci", {"RUN_ALL": "1", "NIGHTLY": "1"}),
(COMMAND_RUN_AMD_CI, "amd-ci", {}),
(COMMAND_RUN_AMD_CI_ALL, "amd-ci", {"RUN_ALL": "1"}),
(
COMMAND_RUN_AMD_CI_NIGHTLY,
"amd-ci",
{"RUN_ALL": "1", "NIGHTLY": "1"},
),
)
for original_command, pipeline, expected_env in cases:
command = f"{original_command} --allow-stale"
for ready, behind in ((False, 100), (True, 0), (True, 100)):
with self.subTest(command=command, ready=ready, behind=behind):
github = FakeGitHub(
permission="read",
behind=behind,
pr=make_pr(
base={"ref": "feature/stack"},
labels=[{"name": "ready"}] if ready else [],
),
)
buildkite = FakeBuildkite([[], []])
run(make_event(command, "author"), github, buildkite)
self.assertEqual(pipeline_for_command(command), pipeline)
if not ready:
self.assertEqual(buildkite.created_builds, [])
self.assertEqual(buildkite.list_calls, [])
self.assertEqual(github.lag_queries, [])
self.assertTrue(github.comments[0].startswith("❌ "))
continue
self.assertEqual(len(buildkite.created_builds), 1)
payload = buildkite.created_builds[0]
self.assertEqual(payload["commit"], github.pr["head"]["sha"])
self.assertEqual(
payload["pull_request_base_branch"], "feature/stack"
)
self.assertEqual(payload["message"], f"PR #42 {command} by @author")
self.assertEqual(
payload["env"],
{
"VLLM_CI_GITHUB_COMMENT_ID": "99",
"VLLM_CI_TRIGGERED_BY": "author",
**expected_env,
},
)
self.assertEqual(
github.lag_queries,
[("feature/stack", github.pr["head"]["sha"])],
)
self.assertIn(
"Buildkite AMD CI #123"
if pipeline == "amd-ci"
else "Buildkite CI #123",
github.comments[0],
)
if behind:
self.assertIn("100 commits behind", github.comments[0])
self.assertIn("upstream `feature/stack`", github.comments[0])
self.assertIn("Before merging", github.comments[0])
self.assertIn("latest PR commit", github.comments[0])
self.assertIn(original_command, github.comments[0])
def test_base_lag_query_qualifies_branch_refs_and_validates_count(self) -> None:
for base_ref, encoded_ref in (
("main", "main"),
("release/1.0", "release%2F1.0"),
("feature/stack", "feature%2Fstack"),
):
for count in (0, 5, None, -1, True, "5"):
with self.subTest(base_ref=base_ref, count=count):
transport = FakeTransport(
{} if count is None else {"behind_by": count}
)
github = GitHubClient("token", "vllm-project/vllm", transport)
if type(count) is int and count >= 0:
self.assertEqual(
github.get_commits_behind_base(base_ref, "head-sha"), count
)
else:
with self.assertRaises(ApiError):
github.get_commits_behind_base(base_ref, "head-sha")
self.assertEqual(
transport.calls[0]["url"],
"https://api.github.com/repos/vllm-project/vllm/compare/"
f"refs%2Fheads%2F{encoded_ref}...head-sha?per_page=1",
)
self.assertEqual(transport.calls[0]["method"], "GET")
def test_unknown_lag_or_changed_pr_prevents_runs(self) -> None:
for command in (
COMMAND_RUN_CI,
COMMAND_RUN_AMD_CI,
f"{COMMAND_RUN_CI} --allow-stale",
f"{COMMAND_RUN_AMD_CI} --allow-stale",
):
for fault in ("api", "head", "closed", "base"):
with self.subTest(command=command, fault=fault):
github = FakeGitHub()
buildkite = FakeBuildkite([[], []])
def check(
base_ref: str,
sha: str,
fault: str = fault,
github: FakeGitHub = github,
) -> int:
if fault == "api":
raise ApiError(503, "Unavailable")
changed = make_pr()
if fault == "head":
changed["head"]["sha"] = "another-head"
elif fault == "closed":
changed["state"] = "closed"
else:
changed["base"]["ref"] = "another-base"
github.prs[42] = changed
return 0
with patch.object(
github, "get_commits_behind_base", side_effect=check
):
run(make_event(command), github, buildkite)
self.assertEqual(buildkite.created_builds, [])
self.assertEqual(buildkite.retry_calls, [])
self.assertTrue(github.comments[0].startswith("❌ "))
self.assertNotIn("Triggered", github.comments[0])
def test_amd_run_ignores_blocked_builds(self) -> None:
for metadata in ({}, {"github-comment-id": "98"}):
with self.subTest(metadata=metadata):
@@ -627,6 +814,7 @@ class RunCiCommandTest(unittest.TestCase):
run(make_event(COMMAND_RUN_CI, "author"), github, buildkite)
self.assertEqual(buildkite.list_calls, [])
self.assertEqual(github.lag_queries, [])
self.assertEqual(github.reactions, ["eyes"])
self.assertTrue(github.comments[0].startswith("❌ "))
self.assertIn("approve the PR", github.comments[0])
@@ -848,6 +1036,7 @@ class RunCiCommandTest(unittest.TestCase):
def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None:
github = FakeGitHub(
behind=100,
permission="read",
pr=make_pr(labels=[{"name": "ready"}]),
)
@@ -867,10 +1056,12 @@ class RunCiCommandTest(unittest.TestCase):
run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite)
self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)])
self.assertEqual(github.lag_queries, [])
self.assertIn("Queued 3 failed job", github.comments[0])
def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None:
github = FakeGitHub(
behind=100,
permission="read",
pr=make_pr(labels=[{"name": "ready"}]),
)
@@ -891,6 +1082,7 @@ class RunCiCommandTest(unittest.TestCase):
run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite)
self.assertEqual(buildkite.retry_calls, [(321, RETRY_STATES)])
self.assertEqual(github.lag_queries, [])
self.assertIn("Buildkite AMD CI #321", github.comments[0])
def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None:
@@ -912,6 +1104,7 @@ class RunCiCommandTest(unittest.TestCase):
def test_ci_retry_creates_filtered_build_for_new_head(self) -> None:
github = FakeGitHub(
behind=100,
permission="read",
pr=make_pr(labels=[{"name": "ready"}]),
)
@@ -953,6 +1146,7 @@ class RunCiCommandTest(unittest.TestCase):
self.assertEqual(len(buildkite.created_builds), 1)
payload = buildkite.created_builds[0]
self.assertEqual(payload["commit"], "0123456789abcdef")
self.assertEqual(github.lag_queries, [])
self.assertEqual(payload["message"], "PR #42 /ci retry by @author")
self.assertEqual(
json.loads(payload["env"]["VLLM_CI_ONLY_STEP_KEYS"]),
@@ -1104,6 +1298,7 @@ class RunCiCommandTest(unittest.TestCase):
run(make_event(COMMAND_CANCEL_CI), github, buildkite)
self.assertEqual(buildkite.cancel_calls, [123, 124, 126])
self.assertEqual(github.lag_queries, [])
self.assertEqual(
[request["branch"] for request in buildkite.list_requests],
["contributor:feature", "feature"],
@@ -1152,6 +1347,7 @@ class RunCiCommandTest(unittest.TestCase):
run(make_event(COMMAND_CANCEL_AMD_CI), github, buildkite)
self.assertEqual(buildkite.cancel_calls, [322, 321])
self.assertEqual(github.lag_queries, [])
self.assertEqual(
[request["branch"] for request in buildkite.list_requests],
["contributor:feature", "feature"],
benchmarks/kernels/benchmark_gelu_and_mul_sparse.py
+90
View File
@@ -0,0 +1,90 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
import torch
import vllm.kernels # noqa: F401
from vllm import ir
from vllm.benchmarks.lib.utils import default_vllm_config
from vllm.triton_utils import triton
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048]
INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768]
CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES))
STD_MULTIPLIER = 1.6448536269514722
@default_vllm_config()
def benchmark_gelu_and_mul_sparse(
num_tokens: int,
intermediate_size: int,
provider: str,
dtype: torch.dtype,
):
set_random_seed(0)
torch.set_default_device("cuda")
x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype)
native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn
args = (x, STD_MULTIPLIER, "tanh")
expected = native(*args)
actual = triton_impl(*args)
torch.testing.assert_close(
actual,
expected,
**ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype),
)
if provider == "triton":
fn = lambda: triton_impl(*args)
elif provider == "native_compiled":
compiled = torch.compile(native, fullgraph=True)
compiled(*args)
fn = lambda: compiled(*args)
else:
fn = lambda: native(*args)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
fn, quantiles=[0.5, 0.2, 0.8]
)
return ms, max_ms, min_ms
if __name__ == "__main__":
parser = FlexibleArgumentParser(
description="Benchmark the Gemma3n sparse GELU activation."
)
parser.add_argument(
"--dtype",
type=str,
choices=["half", "bfloat16", "float"],
default="bfloat16",
)
args = parser.parse_args()
dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
perf_report = triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens", "intermediate_size"],
x_vals=CONFIGS,
line_arg="provider",
line_vals=["triton", "native_compiled", "native_eager"],
line_names=["Triton", "Native (torch.compile)", "Native (eager)"],
styles=[("blue", "-"), ("green", "-"), ("red", "-")],
ylabel="ms",
plot_name="gelu-and-mul-sparse-performance",
args={},
)
)
perf_report(
lambda num_tokens, intermediate_size, provider: (
benchmark_gelu_and_mul_sparse(
num_tokens, intermediate_size, provider, dtype
)
)
).run(print_data=True)
benchmarks/kernels/benchmark_persistent_topk.py
+134
View File
@@ -0,0 +1,134 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows.
Run on a reserved GPU, for example:
chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py
Use --capacity to emulate graphs captured at a larger maximum context length.
Run the same command on the base and candidate builds to compare JSON results.
"""
import argparse
import itertools
import json
import statistics
from pathlib import Path
import torch
from flashinfer.testing import bench_gpu_time_with_cupti
import vllm._custom_ops # noqa: F401
from vllm.platforms import current_platform
def benchmark(rows: int, width: int, k: int, args) -> dict:
capacity = args.capacity or width
if capacity < width:
raise ValueError("Capacity must be at least the valid-length upper bound")
torch.manual_seed(args.seed)
logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32)
if args.full_length:
lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32)
else:
lengths = torch.randint(
int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32
)
invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None]
logits.masked_fill_(invalid, float("nan"))
out = torch.empty((rows, k), device="cuda", dtype=torch.int32)
workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8)
cooperative = (
args.backend == "auto"
and rows <= 64
and capacity % 4 == 0
and current_platform.has_device_capability(90)
and not current_platform.is_device_capability_family(120)
)
op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk
def run():
op(logits, lengths, out, workspace, k, capacity)
run()
valid = torch.arange(k, device="cuda")[None] < lengths[:, None]
assert torch.all(out[~valid] == -1)
assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid)
sorted_indices = out.sort(dim=1).values
assert torch.all(
(sorted_indices[:, 1:] != sorted_indices[:, :-1])
| (sorted_indices[:, 1:] == -1)
)
selected = logits.gather(1, out.clamp_min(0).long()).masked_fill(
~valid, -float("inf")
)
reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values
torch.testing.assert_close(
selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0
)
trials = [
statistics.median(
bench_gpu_time_with_cupti(
run,
use_cuda_graph=True,
cold_l2_cache=not args.warm,
dry_run_iters=5,
repeat_iters=args.repeat,
)
)
* 1000
for _ in range(args.trials)
]
us = statistics.median(trials)
return dict(
rows=rows,
width=width,
capacity=capacity,
k=k,
backend="cooperative" if cooperative else "persistent",
us=us,
trials_us=trials,
# Minimum useful traffic: one valid-score read and one index write.
effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000),
)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024])
parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840])
parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048])
parser.add_argument("--capacity", type=int)
parser.add_argument(
"--full-length", action="store_true", help="Use exactly --widths valid scores"
)
parser.add_argument("--backend", choices=["auto", "persistent"], default="auto")
parser.add_argument(
"--warm", action="store_true", help="Keep L2 warm between replays"
)
parser.add_argument("--repeat", type=int, default=60)
parser.add_argument("--trials", type=int, default=3)
parser.add_argument("--seed", type=int, default=123)
parser.add_argument("--output", type=Path)
args = parser.parse_args()
metadata = dict(
gpu=str(torch.cuda.get_device_properties(0)),
torch=torch.__version__,
cuda=torch.version.cuda,
args=vars(args),
)
print(json.dumps(metadata, default=str))
results = []
for rows, width, k in itertools.product(args.rows, args.widths, args.top_k):
result = benchmark(rows, width, k, args)
results.append(result)
print(json.dumps(result), flush=True)
if args.output:
args.output.write_text(
json.dumps(dict(metadata=metadata, results=results), indent=2, default=str)
+ "\n"
)
if __name__ == "__main__":
main()
benchmarks/multi_turn/benchmark_serving_multi_turn.py
+4 -1
View File
@@ -32,8 +32,11 @@ from bench_dataset import (
from bench_utils import TEXT_SEPARATOR, Color, logger
from transformers import AutoTokenizer # type: ignore
from vllm.benchmarks.lib.utils import redact_sensitive_namespace
NUM_TOKENS_FROM_DATASET = 0
TERM_SIGNAL = None
_SENSITIVE_ARG_FIELDS = ("api_key", "header")
class ConversationSampling(str, Enum):
@@ -1548,7 +1551,7 @@ async def main() -> None:
args = parser.parse_args()
logger.info(args)
logger.info(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS))
logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}")
logger.info(f"url={args.url}")
cmake/external_projects/deepgemm.cmake
+10 -6
View File
@@ -28,9 +28,11 @@ if(DEEPGEMM_SRC_DIR)
message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
else()
# Keep in sync with tools/install_deepgemm.sh
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
# Pinned to the tip of the nv_dev branch (SM120 support).
set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a")
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
# Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120
# port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU
# alpha/beta, and SiTU for FP8/FP4 Mega MoE.
set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4")
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _deepgemm_fc_root)
@@ -51,7 +53,7 @@ else()
BINARY_DIR "${_deepgemm_bin}"
GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit"
GIT_PROGRESS TRUE
)
endif()
@@ -119,13 +121,15 @@ if(DEEPGEMM_ARCHS)
message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
# add_custom_command does no implicit header scanning; glob explicitly so
# header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
# header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild.
file(GLOB_RECURSE _dg_headers
"${deepgemm_SOURCE_DIR}/csrc/*.h"
"${deepgemm_SOURCE_DIR}/csrc/*.hpp"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh"
"${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h"
"${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp")
set(_dg_markers)
set(_dg_seen_soabis)
cmake/external_projects/flashmla.cmake
+59 -26
View File
@@ -19,7 +19,7 @@ else()
FetchContent_Declare(
flashmla
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183
GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS)
set(FlashMLA_SOURCES
${flashmla_SOURCE_DIR}/csrc/api/api.cpp
${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp
${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp
${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp
${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
# Misc kernels for decoding
${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu
# sm90 dense decode
${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
# sm90 sparse decode
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
# sm90 sparse prefill
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
# sm100 dense prefill & backward
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
# sm100 dense prefill (inference-only: dense backward stays out)
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
# sm100 sparse prefill
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
# sm100 fused norm + rope + attn + rope + cast
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
# sm100 sparse decode
${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
)
set(FlashMLA_Extension_SOURCES
@@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS)
set(FlashMLA_INCLUDES
${flashmla_SOURCE_DIR}/csrc
${flashmla_SOURCE_DIR}/csrc/kerutils/include
${flashmla_SOURCE_DIR}/csrc/sm90
${flashmla_SOURCE_DIR}/csrc/cutlass/include
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
)
cmake/external_projects/fmha_sm100.cmake
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa
GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+150 -31
View File
@@ -23,9 +23,15 @@
* cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
* second along last dim; each half is rope_dim/2 = 32 values)
*
* Cache layout per paged-cache block (block_size tokens):
* [0, bs*576): token data, 448 fp8 + 128 bf16 each
* [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
* Cache layout per paged-cache block (block_size tokens), selected by the
* `kv_mxfp8` argument:
* V4 (584 B/token, kv_mxfp8=false)
* [0, bs*576): token data, 448 fp8 + 128 bf16 each
* [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
* V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the
* whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims)
* [0, bs*512): token data, 512 fp8 each
* [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token
*/
#include "torch_utils.h"
@@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64;
constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims.
constexpr int kMxfp8QuantBlock = 32;
constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16
constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512
// FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw
// 240.0) to match the rest of vLLM's FNUZ pipeline.
#if defined(USE_ROCM) && defined(__gfx942__)
@@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
// ────────────────────────────────────────────────────────────────────────────
// Small inline helpers
// ────────────────────────────────────────────────────────────────────────────
__device__ __forceinline__ float warp2MaxAbs(float val) {
// Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block).
return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1));
}
__device__ __forceinline__ float warp4MaxAbs(float val) {
// Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
@@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) {
// slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant
// + paged-cache
// insert)
//
// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a
// runtime argument rather than doubling every (head count, q-norm)
// instantiation of a kernel whose hot path is the Q branch.
template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
__device__ __forceinline__ void processDeepseekV4Slot(
uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
@@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
int64_t const* __restrict__ position_ids,
float const* __restrict__ cos_sin_cache, int const cache_block_size,
int const kv_block_stride) {
int const kv_block_stride, bool const kv_mxfp8) {
using Converter = vllm::_typeConvert<scalar_t_in>;
bool const isKV = (slotIdx == kNumHeadsQPadded);
bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
@@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot(
*reinterpret_cast<uint4*>(dst) = out0;
*reinterpret_cast<uint4*>(dst + 8) = out1;
} else {
// ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
// ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16;
// V4.1 (kv_mxfp8) quantizes them too.
int64_t const slot_id = slot_mapping[tokenIdx];
if (slot_id >= 0) {
int const token_data_bytes =
kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes;
int const scale_bytes =
kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken;
int64_t const block_idx = slot_id / cache_block_size;
int64_t const pos_in_block = slot_id % cache_block_size;
uint8_t* block_base =
k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes;
uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
uint8_t* token_scale_ptr =
block_base +
static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
pos_in_block * kScaleBytesPerToken;
static_cast<int64_t>(cache_block_size) * token_data_bytes +
pos_in_block * scale_bytes;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
@@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot(
for (int i = 0; i < kElemsPerLane; i++) {
local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
}
float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
float const absmax = fmaxf(
kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax),
1e-4f);
float const exponent = ceilf(log2f(absmax / kFp8Max));
float const inv_scale = exp2f(-exponent);
if (!is_rope_lane) {
if (kv_mxfp8 || !is_rope_lane) {
uint8_t out_bytes[kElemsPerLane];
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
@@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot(
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
*reinterpret_cast<uint4 const*>(out_bytes);
if ((laneId & 3) == 0) {
int const q_block_idx = laneId >> 2;
float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
}
if (laneId == 0) {
token_scale_ptr[kNumQuantBlocks] = 0;
float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
if (kv_mxfp8) {
// One scale per 32 dims = one per lane pair.
if ((laneId & 1) == 0) {
token_scale_ptr[laneId >> 1] = static_cast<uint8_t>(encoded);
}
} else {
if ((laneId & 3) == 0) {
token_scale_ptr[laneId >> 2] = static_cast<uint8_t>(encoded);
}
if (laneId == 0) {
token_scale_ptr[kNumQuantBlocks] = 0;
}
}
} else {
uint4 out0, out1;
@@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
int const num_heads_q, // live Q heads (input layout)
int const cache_block_size, // tokens per paged-cache block
int const kv_block_stride) { // bytes per paged-cache block
int const kv_block_stride, // bytes per paged-cache block
bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
// BF16 _typeConvert specialization is unavailable on pre-Ampere. The
// DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
@@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
kv_block_stride);
kv_block_stride, kv_mxfp8);
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
@@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
float const* __restrict__ cos_sin_cache, float const eps,
int const num_tokens_full, int const num_tokens_insert,
int const num_heads_q, int const cache_block_size,
int const kv_block_stride) {
int const kv_block_stride, bool const kv_mxfp8) {
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
return;
@@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
cos_sin_cache, cache_block_size, kv_block_stride);
cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8);
// ── Buffer rotation: hand the prefetched LDGs to the next iter.
v0_curr = v0_next;
@@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated(
uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
float const* cos_sin_cache, float const eps, int const num_tokens_full,
int const num_tokens_insert, int const num_heads_q,
int const cache_block_size, int const kv_block_stride,
int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8,
cudaStream_t stream) {
constexpr int kBlockSize = 256;
constexpr int kWarpsPerBlock = kBlockSize / 32;
@@ -624,13 +657,19 @@ static void launchFusedDeepseekV4Templated(
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
if (num_tokens_full < NUM_TOKEN_CUTOFF) {
// The reduced grid pays off by iterating a token's Q heads inside one CTA.
// With zero query heads there is nothing to iterate, so one warp per block
// would do all the work and the other seven idle; keep the warp-per-slot
// grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single
// insert, where it is worth ~2x.
if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) {
cudaLaunchKernelEx(&config,
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<
scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
cos_sin_cache, eps, num_tokens_full, num_tokens_insert,
num_heads_q, cache_block_size, kv_block_stride);
num_heads_q, cache_block_size, kv_block_stride,
kv_mxfp8);
} else {
config.gridDim = dim3(num_tokens_full);
cudaLaunchKernelEx(
@@ -639,7 +678,7 @@ static void launchFusedDeepseekV4Templated(
scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
kv_block_stride);
kv_block_stride, kv_mxfp8);
}
#else
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
@@ -649,7 +688,7 @@ static void launchFusedDeepseekV4Templated(
<<<grid, kBlockSize, 0, stream>>>(
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
cache_block_size, kv_block_stride);
cache_block_size, kv_block_stride, kv_mxfp8);
#endif
}
@@ -663,19 +702,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
int const num_tokens_full, int const num_tokens_insert,
int const num_heads_q, int const num_heads_q_padded,
int const cache_block_size, int const kv_block_stride,
bool const apply_q_norm, cudaStream_t stream) {
bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) {
#define DISPATCH(N) \
case N: \
if (apply_q_norm) { \
launchFusedDeepseekV4Templated<scalar_t_in, N, true>( \
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
num_heads_q, cache_block_size, kv_block_stride, stream); \
num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
} else { \
launchFusedDeepseekV4Templated<scalar_t_in, N, false>( \
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
num_heads_q, cache_block_size, kv_block_stride, stream); \
num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
} \
return;
@@ -959,6 +998,86 @@ static void launchFullCacheKernel(
// ────────────────────────────────────────────────────────────────────────────
// Torch op wrapper
// ────────────────────────────────────────────────────────────────────────────
void fused_deepseek_v4_kv_rope_insert(
torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
torch::stable::Tensor const& slot_mapping,
torch::stable::Tensor const& position_ids,
torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8) {
using torch::headeronly::ScalarType;
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() &&
kv.dim() == 2 && kv.size(1) == 512 &&
kv.scalar_type() == ScalarType::BFloat16,
"kv must be contiguous CUDA bfloat16 [tokens, 512]");
auto const device = kv.get_device_index();
STD_TORCH_CHECK(k_cache.device().is_cuda() &&
k_cache.get_device_index() == device &&
k_cache.dim() == 3 && cache_block_size > 0 &&
k_cache.size(1) == cache_block_size &&
k_cache.stride(2) == 1,
"k_cache must be a paged cache on the KV device");
for (auto const* tensor : {&slot_mapping, &position_ids}) {
STD_TORCH_CHECK(tensor->device().is_cuda() &&
tensor->get_device_index() == device &&
tensor->scalar_type() == ScalarType::Long &&
tensor->dim() == 1 && tensor->is_contiguous(),
"slots and positions must be contiguous int64 on the KV device");
}
STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) &&
slot_mapping.size(0) <= kv.size(0),
"positions must match KV rows; slots may omit padded rows");
STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
cos_sin_cache.get_device_index() == device &&
cos_sin_cache.scalar_type() == ScalarType::Float &&
cos_sin_cache.dim() == 2 &&
cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(),
"cos_sin_cache must be contiguous float32 [max_pos, 64]");
bool const packed = k_cache.scalar_type() == ScalarType::Byte;
bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn;
STD_TORCH_CHECK(packed || fp8 ||
k_cache.scalar_type() == ScalarType::BFloat16,
"cache must be uint8, bfloat16 or float8_e4m3fn");
STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512),
"cache row width does not match its dtype");
if (fp8) {
STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() &&
fp8_scale->get_device_index() == device &&
fp8_scale->scalar_type() == ScalarType::Float &&
fp8_scale->numel() == 1,
"fp8 cache requires one float32 KV scale on the KV device");
}
int const num_tokens = static_cast<int>(slot_mapping.size(0));
if (num_tokens == 0) return;
const torch::stable::accelerator::DeviceGuard device_guard(device);
const cudaStream_t stream = get_current_cuda_stream(device);
using scalar_t = c10::BFloat16;
auto const* input = reinterpret_cast<scalar_t const*>(kv.const_data_ptr());
auto* cache = reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr());
auto const* slots = slot_mapping.const_data_ptr<int64_t>();
auto const* positions = position_ids.const_data_ptr<int64_t>();
auto const* cos_sin = cos_sin_cache.const_data_ptr<float>();
// Zero query heads schedules only the existing KV branch, preserving its
// RoPE rounding and quantization without allocating any query tensors.
if (packed) {
vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated<scalar_t, 0, false>(
nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f,
num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
static_cast<int>(k_cache.stride(0)), kv_mxfp8, stream);
} else if (fp8) {
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, true>(
nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin,
fp8_scale->const_data_ptr<float>(), nullptr, 0.0f, num_tokens, num_tokens,
0, static_cast<int>(cache_block_size), k_cache.stride(0),
k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream);
} else {
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false, false>(
nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr,
nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast<int>(cache_block_size),
k_cache.stride(0) * 2, k_cache.stride(1) * 2, false,
"fused_deepseek_v4_kv_rope_insert", stream);
}
}
torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16
torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only)
@@ -967,7 +1086,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
torch::stable::Tensor const& position_ids, // [N] int64
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
int64_t q_head_padded, // padded Q head count for output
double eps, int64_t cache_block_size, bool apply_q_norm) {
double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) {
STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
"q_in must be contiguous CUDA");
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
@@ -1036,7 +1155,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
num_tokens_full, num_tokens_insert, num_heads_q,
num_heads_q_padded, cache_block_size_i, kv_block_stride,
apply_q_norm, stream);
apply_q_norm, kv_mxfp8, stream);
});
return q_out;
}
csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu
+2 -12
View File
@@ -109,9 +109,6 @@ __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw,
#endif
}
// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload
// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which
// would saturate the tile's largest values.
__device__ __forceinline__ float nvfp4_tile_scale(const float* vals,
float max_val,
uint8_t* sf_out) {
@@ -121,15 +118,8 @@ __device__ __forceinline__ float nvfp4_tile_scale(const float* vals,
amax = fmaxf(amax, fabsf(vals[i]));
}
const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9
__nv_fp8_e4m3 sf8(sf_f);
uint8_t sf_bits = *reinterpret_cast<const uint8_t*>(&sf8);
if (float(sf8) < sf_f && sf_bits < 0x7E) {
// Bump to the next e4m3 value (bit patterns of positive e4m3 values are
// monotonic; 0x7E = 448 is the max finite value).
sf_bits += 1;
sf8 = *reinterpret_cast<const __nv_fp8_e4m3*>(&sf_bits);
}
*sf_out = sf_bits;
const __nv_fp8_e4m3 sf8(sf_f);
*sf_out = *reinterpret_cast<const uint8_t*>(&sf8);
return float(sf8);
}
csrc/libtorch_stable/ops.h
+8 -1
View File
@@ -259,12 +259,19 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
torch::stable::Tensor& position_ids,
int64_t forced_token_heads_per_warp);
void fused_deepseek_v4_kv_rope_insert(
torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
torch::stable::Tensor const& slot_mapping,
torch::stable::Tensor const& position_ids,
torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8);
torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
torch::stable::Tensor const& position_ids,
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
double eps, int64_t cache_block_size, bool apply_q_norm);
double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8);
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
csrc/libtorch_stable/persistent_topk.cuh
+52 -34
View File
@@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB
/*!
* \brief Filtered Top-K kernel for ragged sequences.
*
* \tparam DType Data type (float, half, nv_bfloat16)
* \tparam IdType Index type (int32_t)
* \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
*/
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
bool UsePredicatedShortLoads = false>
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
IdType* __restrict__ output,
const IdType* __restrict__ lengths,
uint32_t num_rows, uint32_t top_k,
uint32_t max_len) {
template <uint32_t MAX_K>
struct FilteredTopKStorage {
alignas(128) int histogram[2][256 + 128];
alignas(128) int counter;
alignas(128) int threshold_bin;
alignas(128) int num_input[2];
alignas(128) int indices[MAX_K];
int last_remain;
};
// With CheckOverflow, return false before using a truncated stash so the caller
// can retry with exact full-row selection.
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K,
bool UsePredicatedShortLoads, bool CheckOverflow = false>
__device__ __forceinline__ bool filtered_topk_row(
const DType* score, IdType* dst, int length, uint32_t top_k,
FilteredTopKStorage<MAX_K>& storage) {
constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS;
constexpr int RADIX = 256;
constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE;
const uint32_t bid = blockIdx.x;
const int tx = threadIdx.x;
if (bid >= num_rows) return;
const int length =
(lengths != nullptr) ? lengths[bid] : static_cast<int>(max_len);
const DType* score = input + bid * max_len;
IdType* dst = output + bid * top_k;
// Trivial case: length <= top_k
if (length <= static_cast<int>(top_k)) {
for (int i = tx; i < static_cast<int>(top_k); i += BLOCK_SIZE) {
dst[i] = (i < length) ? static_cast<IdType>(i) : static_cast<IdType>(-1);
}
return;
return true;
}
// Short path
@@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
_smem_reg);
}
return;
return true;
}
// Static shared memory
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
alignas(128) __shared__ int s_counter;
alignas(128) __shared__ int s_threshold_bin_id;
alignas(128) __shared__ int s_num_input[2];
alignas(128) __shared__ int s_indices[MAX_K];
auto& s_histogram_buf = storage.histogram;
auto& s_counter = storage.counter;
auto& s_threshold_bin_id = storage.threshold_bin;
auto& s_num_input = storage.num_input;
auto& s_indices = storage.indices;
auto& s_histogram = s_histogram_buf[0];
// Dynamic shared memory for input double buffer
@@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
// Stage 2: refine with 8bit radix passes
#pragma unroll
for (int round = 0; round < NUM_ROUNDS; ++round) {
__shared__ int s_last_remain;
auto& s_last_remain = storage.last_remain;
const auto r_idx = round % 2;
const auto _raw_num_input = s_num_input[r_idx];
if constexpr (CheckOverflow) {
if (_raw_num_input > SMEM_INPUT_SIZE) return false;
}
const auto num_input =
(_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE;
@@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
const int idx = s_indices[base];
dst[base] = static_cast<IdType>(idx);
}
return true;
}
/*!
* \brief Filtered Top-K kernel for ragged sequences.
*
* \tparam DType Data type (float, half, nv_bfloat16)
* \tparam IdType Index type (int32_t)
* \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
*/
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
bool UsePredicatedShortLoads = false>
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
IdType* __restrict__ output,
const IdType* __restrict__ lengths,
uint32_t num_rows, uint32_t top_k,
uint32_t max_len) {
const uint32_t bid = blockIdx.x;
if (bid >= num_rows) return;
const int length = lengths ? lengths[bid] : static_cast<int>(max_len);
__shared__ FilteredTopKStorage<MAX_K> storage;
filtered_topk_row<DType, IdType, VEC_SIZE, MAX_K, UsePredicatedShortLoads>(
input + bid * max_len, output + bid * top_k, length, top_k, storage);
}
// Helper to compute GCD for VEC_SIZE selection
csrc/libtorch_stable/sampled_topk.cuh
+253
View File
@@ -0,0 +1,253 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect:
// https://github.com/deepseek-ai/DeepSelect
// A sample estimates a cutoff, then survivors are compacted in shared memory
// for exact FP32 selection. Too few or too many candidates trigger an exact
// full-row fallback.
#pragma once
#include <cub/cub.cuh>
#include <cuda_runtime.h>
#include <cstdint>
#include "persistent_topk.cuh"
namespace vllm::sampled_topk {
constexpr int kThreads = 1024;
// Conservative crossover bounds from the B300 batch/length sweep.
template <int K>
constexpr int kMinSampledLength = K == 512 ? 98304 : 65536;
// Target half the buffer to leave room for sampling error.
constexpr int kSample = 4096;
constexpr int kCapacity = 8192;
struct Storage {
uint2 candidates[kCapacity];
int histogram[2048];
typename cub::BlockScan<int, kThreads>::TempStorage scan;
int count;
int remaining;
int bin;
uint32_t prefix;
};
__device__ __forceinline__ uint32_t ordered(float value) {
const uint32_t bits = __float_as_uint(value);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
__device__ __forceinline__ uint32_t coarse(uint32_t key) {
const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key;
return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits));
}
// Select the bin containing the rank and update the rank within that bin.
template <int Bins>
__device__ void select_bin(Storage& s) {
constexpr int kItems = (Bins + kThreads - 1) / kThreads;
int sums[kItems];
#pragma unroll
for (int j = 0; j < kItems; ++j) {
const int bin = Bins - 1 - (threadIdx.x * kItems + j);
sums[j] = bin >= 0 ? s.histogram[bin] : 0;
}
cub::BlockScan<int, kThreads>(s.scan).InclusiveSum(sums, sums);
const int remaining = s.remaining;
__syncthreads();
#pragma unroll
for (int j = 0; j < kItems; ++j) {
const int bin = Bins - 1 - (threadIdx.x * kItems + j);
if (bin >= 0 && sums[j] >= remaining &&
sums[j] - s.histogram[bin] < remaining) {
s.bin = bin;
s.remaining = remaining - (sums[j] - s.histogram[bin]);
}
}
__syncthreads();
}
// Find a coarse bin, then refine its full FP32 keys one byte at a time.
// Read shared candidates normally, or the full row for exact fallback.
template <bool Buffered>
__device__ uint32_t threshold(Storage& s, const float* row, int count,
int rank) {
const int tid = threadIdx.x;
if (tid == 0) {
s.remaining = rank;
s.prefix = 0;
}
for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
__syncthreads();
for (int i = tid; i < count; i += kThreads) {
const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
atomicAdd(&s.histogram[coarse(key)], 1);
}
__syncthreads();
select_bin<2048>(s);
const int coarse_bin = s.bin;
uint32_t mask = 0;
for (int shift = 24; shift >= 0; shift -= 8) {
if (tid < 256) s.histogram[tid] = 0;
__syncthreads();
const uint32_t prefix = s.prefix;
for (int i = tid; i < count; i += kThreads) {
const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
if (coarse(key) == coarse_bin && (key & mask) == prefix) {
atomicAdd(&s.histogram[(key >> shift) & 255], 1);
}
}
__syncthreads();
select_bin<256>(s);
if (tid == 0) s.prefix = prefix | (static_cast<uint32_t>(s.bin) << shift);
mask |= 255u << shift;
__syncthreads();
}
return s.prefix;
}
__device__ __forceinline__ int reserve(bool hit, int* counter) {
const uint32_t mask = __ballot_sync(0xffffffffu, hit);
if (mask == 0) return 0;
const int lane = threadIdx.x & 31;
int base = 0;
if (lane == 0) base = atomicAdd(counter, __popc(mask));
base = __shfl_sync(0xffffffffu, base, 0);
return base + __popc(mask & ((1u << lane) - 1));
}
template <bool Buffered>
__device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
uint32_t cutoff) {
const int tid = threadIdx.x;
if (tid == 0) s.count = 0;
__syncthreads();
for (int base = 0; base < count; base += kThreads) {
const int i = base + tid;
const uint2 pair = i < count ? (Buffered ? s.candidates[i]
: make_uint2(i, ordered(row[i])))
: make_uint2(0, 0);
bool keep = i < count && pair.y > cutoff;
if (i < count && pair.y == cutoff) {
keep = atomicSub(&s.remaining, 1) > 0;
}
const int offset = reserve(keep, &s.count);
if (keep) dst[offset] = pair.x;
}
}
template <int K>
__global__ void __launch_bounds__(kThreads)
sampled_topk_kernel(const float* __restrict__ input,
const int32_t* __restrict__ lengths,
int32_t* __restrict__ output, int64_t stride,
int max_length) {
extern __shared__ __align__(16) unsigned char smem[];
auto& s = *reinterpret_cast<Storage*>(smem);
const int tid = threadIdx.x;
const int length = max(0, min(lengths[blockIdx.x], max_length));
const float* row = input + blockIdx.x * stride;
int32_t* dst = output + blockIdx.x * K;
// Graph allocations can hold rows much shorter than the host length bound.
if (length < kMinSampledLength<K>) {
__shared__ filtered_topk::FilteredTopKStorage<K> fallback_storage;
bool selected;
if ((stride & 3) == 0 && (reinterpret_cast<uintptr_t>(row) & 15) == 0) {
selected =
filtered_topk::filtered_topk_row<float, int32_t, 4, K, false, true>(
row, dst, length, K, fallback_storage);
} else {
selected =
filtered_topk::filtered_topk_row<float, int32_t, 1, K, true, true>(
row, dst, length, K, fallback_storage);
}
if (!selected) {
const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
emit<false>(s, row, dst, length, exact_cutoff);
}
return;
}
if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length);
for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
__syncthreads();
// Sample contiguous warps spread across the row to keep loads coalesced.
for (int i = tid; i < kSample; i += kThreads) {
const int index = (i / 32) * (length / (kSample / 32)) + i % 32;
atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1);
}
__syncthreads();
select_bin<2048>(s);
const uint16_t half_key = s.bin << 5;
const uint16_t half_bits =
(half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key;
const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits)));
if (tid == 0) s.count = 0;
__syncthreads();
constexpr int kItems = 16;
// Compact survivors with one reservation per warp per tile. Keep counting
// beyond capacity so overflow triggers fallback instead of using a subset.
for (int base = 0; base < length; base += kItems * kThreads) {
uint32_t keys[kItems];
uint32_t hit_mask = 0;
if ((reinterpret_cast<uintptr_t>(row) & 15) == 0 &&
base + kItems * kThreads <= length) {
#pragma unroll
for (int j = 0; j < kItems / 4; ++j) {
const float4 values =
reinterpret_cast<const float4*>(row + base)[tid + j * kThreads];
keys[4 * j] = ordered(values.x);
keys[4 * j + 1] = ordered(values.y);
keys[4 * j + 2] = ordered(values.z);
keys[4 * j + 3] = ordered(values.w);
}
} else {
#pragma unroll
for (int j = 0; j < kItems; ++j) {
const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
keys[j] = i < length ? ordered(row[i]) : 0;
}
}
#pragma unroll
for (int j = 0; j < kItems; ++j) {
const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j;
}
const int lane = tid & 31;
const int hits = __popc(hit_mask);
int inclusive = hits;
#pragma unroll
for (int delta = 1; delta < 32; delta *= 2) {
const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta);
if (lane >= delta) inclusive += previous;
}
const int total = __shfl_sync(0xffffffffu, inclusive, 31);
int warp_base = 0;
if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total);
warp_base = __shfl_sync(0xffffffffu, warp_base, 31);
int offset = warp_base + inclusive - hits;
#pragma unroll
for (int j = 0; j < kItems; ++j) {
if (hit_mask & (1u << j)) {
if (offset < kCapacity) {
s.candidates[offset] = make_uint2(
base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]);
}
++offset;
}
}
}
__syncthreads();
const int count = s.count;
if (count >= K && count <= kCapacity) {
const uint32_t exact_cutoff = threshold<true>(s, row, count, K);
emit<true>(s, row, dst, count, exact_cutoff);
} else {
const uint32_t exact_cutoff = threshold<false>(s, row, length, K);
emit<false>(s, row, dst, length, exact_cutoff);
}
}
} // namespace vllm::sampled_topk
csrc/libtorch_stable/sampler.cu
+205 -58
View File
@@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
}
template <int step, int kNumThreadsPerBlock, int kNumBins, int kNumFinalItems,
bool multipleBlocksPerRow, bool mergeBlocks, typename SmemFinalType,
bool multipleBlocksPerRow, bool mergeBlocks,
bool useTopK512Optimization, typename SmemFinalType,
typename SmemOutputType>
__device__ bool processHistogramStep(
const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern,
@@ -180,10 +181,15 @@ __device__ bool processHistogramStep(
<< patternShift;
}
int negativeInfinityCount = 0;
auto distributeToBins = [&](float logit, int /* idx */ = 0) {
if (isPartialMatch<patternShift>(logit, logitPattern)) {
uint32_t binIdx = extractBinIdx<step>(logit);
atomicAdd(&smemFinal.histo.data[binIdx], 1);
if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) {
++negativeInfinityCount;
} else {
atomicAdd(&smemFinal.histo.data[binIdx], 1);
}
}
};
@@ -198,6 +204,12 @@ __device__ bool processHistogramStep(
distributeToBins(logit, idx);
}
}
if constexpr (useTopK512Optimization) {
if (negativeInfinityCount > 0) {
const auto bin = extractBinIdx<step>(-INFINITY);
atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount);
}
}
// Make sure the histogram is ready.
__syncthreads();
@@ -254,6 +266,12 @@ __device__ bool processHistogramStep(
// The threshold bin.
thresholdBinIdx = smemThresholdBinIdx[0];
// Resolve a -inf cutoff with exact tie emission, keeping sort padding out.
const bool finalBinFits =
smemFinalBinSize[0] <= kNumFinalItems &&
!(useTopK512Optimization && step < 3 &&
isPartialMatch<patternShift>(-INFINITY, logitPattern) &&
thresholdBinIdx == extractBinIdx<step>(-INFINITY));
auto processBins = [&](float logit, int idx) {
if (isPartialMatch<patternShift>(logit, logitPattern)) {
@@ -262,8 +280,7 @@ __device__ bool processHistogramStep(
// 1. This is step 0 and the threshold bin is small enough (no step 1)
// 2. This is step >= 1 (where pattern matching filters correctly)
// This prevents duplicates when step 0 and step 1 both run.
bool shouldWriteDirectly =
(step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1);
if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
// The element is part of the top-k selection
int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
@@ -279,8 +296,7 @@ __device__ bool processHistogramStep(
}
if constexpr (step < 3) {
// Only fill the final items for sorting if the threshold bin fits
if (binIdx == thresholdBinIdx &&
smemFinalBinSize[0] <= kNumFinalItems) {
if (binIdx == thresholdBinIdx && finalBinFits) {
int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1);
smemFinal.items.logits[dstIdx] = logit;
if constexpr (mergeBlocks) {
@@ -325,20 +341,21 @@ __device__ bool processHistogramStep(
__syncthreads();
// Check if we should continue to next step
return smemFinalBinSize[0] > kNumFinalItems;
return !finalBinFits;
}
// Follows half - 11 - 11 - 10 bit iterations
// Keep the adaptive kernel's device instantiation separate from legacy calls.
template <int kNumThreadsPerBlock, int kNumBins, bool useRadixSort,
bool multipleBlocksPerRow = false, bool mergeBlocks = false,
bool deviceLengthAware = false>
bool deviceLengthAware = false, bool useTopK512Optimization = false>
static __device__ void topKPerRowJob(const int* indices, const float* logits,
int rowStart, int rowEnd, int* outIndices,
float* outLogits, int stride1, int topK) {
static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks);
static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024);
// The number of slots for the final pass.
static constexpr int kNumFinalItems = 2048;
static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048;
// The number of elements per thread for the final sort.
static constexpr int kNumFinalItemsPerThread =
kNumFinalItems / kNumThreadsPerBlock;
@@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
rowIt += kNumThreadsPerBlock) {
outIndices[rowIt] = -1;
if constexpr (multipleBlocksPerRow) {
outLogits[rowIt] = -FLT_MAX;
outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
}
}
@@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
// Step 0: Process first 11 bits of half representation
bool continueToNextStep =
processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
multipleBlocksPerRow, mergeBlocks>(
multipleBlocksPerRow, mergeBlocks,
useTopK512Optimization>(
indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
// Step 1: Process next 11 bits
continueToNextStep =
processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
multipleBlocksPerRow, mergeBlocks>(
multipleBlocksPerRow, mergeBlocks,
useTopK512Optimization>(
indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
// Step 2: Process next 11 bits
continueToNextStep =
processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
multipleBlocksPerRow, mergeBlocks>(
multipleBlocksPerRow, mergeBlocks,
useTopK512Optimization>(
indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
if (continueToNextStep) {
// Step 3: Process last 10 bits
processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems,
multipleBlocksPerRow, mergeBlocks>(
multipleBlocksPerRow, mergeBlocks,
useTopK512Optimization>(
indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput,
smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize,
smemFoundTopKValues, smemFinal, stride1, rowStart, topK);
@@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
// The histogram did not proceed to the final 10 bits, therefore we need to
// sort the final items The logits of the elements to be sorted in the final
// pass.
if constexpr (useRadixSort) {
// Sorting with radix sort
float finalLogits[kNumFinalItemsPerThread];
// The indices of the elements to be sorted in the final pass.
int finalIndices[kNumFinalItemsPerThread];
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
finalLogits[ii] = -FLT_MAX;
}
// Read the elements from SMEM.
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
if (srcIdx < smemFinalDstIdx[0]) {
finalLogits[ii] = smemFinal.items.logits[srcIdx];
finalIndices[ii] = smemFinal.items.indices[srcIdx];
}
}
// Make sure the shared memory has been read.
__syncthreads();
// Sort the elements.
FinalSort(smemFinal.finalSort)
.SortDescendingBlockedToStriped(finalLogits, finalIndices);
// Copy the data back to the shared memory storage.
int baseIdx = smemFoundTopKValues[0];
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
int dstIdx = baseIdx + srcIdx;
if (dstIdx < topK) {
smemOutput[dstIdx] = finalIndices[ii];
if constexpr (multipleBlocksPerRow) {
reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
finalLogits[ii];
}
}
}
} else {
auto insertionSort = [&]() {
// Sorting with insertion sort
auto baseIdx = smemFoundTopKValues[0];
for (int i = threadIdx.x; i < smemFinalDstIdx[0];
@@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits,
}
}
}
};
if constexpr (useRadixSort) {
if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) {
insertionSort();
} else {
// Sorting with radix sort
float finalLogits[kNumFinalItemsPerThread];
// The indices of the elements to be sorted in the final pass.
int finalIndices[kNumFinalItemsPerThread];
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX;
if constexpr (useTopK512Optimization) finalIndices[ii] = -1;
}
// Read the elements from SMEM.
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
if (srcIdx < smemFinalDstIdx[0]) {
finalLogits[ii] = smemFinal.items.logits[srcIdx];
finalIndices[ii] = smemFinal.items.indices[srcIdx];
}
}
// Make sure the shared memory has been read.
__syncthreads();
// Sort the elements.
FinalSort(smemFinal.finalSort)
.SortDescendingBlockedToStriped(finalLogits, finalIndices);
// Copy the data back to the shared memory storage.
int baseIdx = smemFoundTopKValues[0];
#pragma unroll
for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) {
int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x;
int dstIdx = baseIdx + srcIdx;
if (dstIdx < topK) {
smemOutput[dstIdx] = finalIndices[ii];
if constexpr (multipleBlocksPerRow) {
reinterpret_cast<float*>(smemOutput + topK)[dstIdx] =
finalLogits[ii];
}
}
}
}
} else {
insertionSort();
}
__syncthreads();
}
@@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware(
stride1, topK);
}
#ifdef USE_ROCM
// The grid and workspace stay fixed while graph replays change device lengths.
__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows,
int maxBlocks) {
int activeBlocks = 1;
if (numRows <= 32 && rowLength > 16 * 1024) {
activeBlocks = rowLength <= 128 * 1024 ? 8
: rowLength <= 256 * 1024 ? 10
: 16;
} else if (numRows <= 128 && rowLength > 64 * 1024) {
activeBlocks = rowLength <= 512 * 1024 ? 4 : 10;
} else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) {
activeBlocks = numRows <= 160 ? 3 : 2;
}
// Every block in a merged row must contain at least k real indices.
return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512));
}
template <int kNumThreadsPerBlock, bool mergeBlocks = false>
static __global__
__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware(
const float* logits, const int* seqLens, int* outIndices,
int* outIndicesAux, float* outLogitsAux, int stride0, int next_n,
int seqLensIs2D, int maxBlocks) {
#if defined(__gfx950__)
constexpr int topK = 512;
constexpr int kNumBins = 2048;
const int rowIdx = blockIdx.x;
const int rowEnd =
seqLensIs2D
? max(0, seqLens[rowIdx])
: max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1);
const int activeBlocks =
gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks);
outIndices += static_cast<int64_t>(rowIdx) * topK;
if constexpr (mergeBlocks) {
if (activeBlocks == 1) return;
const int64_t offset = static_cast<int64_t>(rowIdx) * maxBlocks * topK;
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, true, false,
true>(outIndicesAux + offset, outLogitsAux + offset, 0,
activeBlocks * topK, outIndices, nullptr, 1, topK);
} else {
if (blockIdx.y >= activeBlocks) return;
logits += static_cast<int64_t>(rowIdx) * stride0;
if (activeBlocks == 1) {
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, false, false, false,
true>(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1,
topK);
return;
}
const int blockSize = rowEnd / activeBlocks;
const int rowStart = blockSize * blockIdx.y;
const int blockRowEnd =
blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize;
const int64_t offset =
(static_cast<int64_t>(rowIdx) * maxBlocks + blockIdx.y) * topK;
topKPerRowJob<kNumThreadsPerBlock, kNumBins, true, true, false, false,
true>(nullptr, logits, rowStart, blockRowEnd,
outIndicesAux + offset, outLogitsAux + offset, 1, topK);
}
#endif
}
#endif
} // namespace vllm
void apply_repetition_penalties_(
@@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
// the same seq_len and the kernel computes the per-row offset itself.
int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0;
#ifdef USE_ROCM
const bool useGfx950DeviceLengthAwareTopK512 =
topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 &&
numColumns <= 1024 * 1024 &&
std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0;
if (useGfx950DeviceLengthAwareTopK512) {
int multipleBlocksPerRowConfig = 1;
if (numRows <= 32 && numColumns > 16 * 1024) {
multipleBlocksPerRowConfig = 16;
} else if (numRows <= 128 && numColumns > 64 * 1024) {
multipleBlocksPerRowConfig = 10;
} else if (numRows <= 160 && numColumns > 256 * 1024) {
multipleBlocksPerRowConfig = 3;
} else if (numRows <= 256 && numColumns > 256 * 1024) {
multipleBlocksPerRowConfig = 2;
}
constexpr int kNumThreadsPerSplitBlock = 1024;
if (multipleBlocksPerRowConfig == 1) {
vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
<<<numRows, kNumThreadsPerSplitBlock, 2 * topK * sizeof(int32_t),
stream>>>(logits.const_data_ptr<float>(),
seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(), nullptr, nullptr,
static_cast<int>(stride0), static_cast<int>(next_n),
seqLensIs2D, multipleBlocksPerRowConfig);
return;
}
const auto outIndicesAux = torch::stable::empty(
{numRows, multipleBlocksPerRowConfig, topK},
torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
const auto outLogitsAux = torch::stable::empty(
{numRows, multipleBlocksPerRowConfig, topK},
torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerSplitBlock>
<<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerSplitBlock,
2 * topK * sizeof(int32_t), stream>>>(
logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(),
outIndicesAux.mutable_data_ptr<int>(),
outLogitsAux.mutable_data_ptr<float>(), static_cast<int>(stride0),
static_cast<int>(next_n), seqLensIs2D, multipleBlocksPerRowConfig);
constexpr int kNumThreadsPerBlockMerge = 1024;
vllm::topKPerRowDecode512DeviceLengthAware<kNumThreadsPerBlockMerge, true>
<<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
nullptr, seqLens.const_data_ptr<int>(),
indices.mutable_data_ptr<int>(),
outIndicesAux.mutable_data_ptr<int>(),
outLogitsAux.mutable_data_ptr<float>(), 0, static_cast<int>(next_n),
seqLensIs2D, multipleBlocksPerRowConfig);
return;
}
#endif
if (numColumns < kSortingAlgorithmThreshold) {
// Use insertion sort
vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
csrc/libtorch_stable/topk.cu
+17 -1
View File
@@ -9,6 +9,7 @@
#ifndef USE_ROCM
#include "persistent_topk.cuh"
#include "sampled_topk.cuh"
#endif
namespace {
@@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
max_smem_per_block = device_prop->sharedMemPerBlockOptin;
}
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
// Allow static fallback storage in addition to the 128 KiB dynamic buffer.
if (num_rows > 64 &&
max_seq_len >= vllm::sampled_topk::kMinSampledLength<TopK> &&
max_smem_per_block >= 144 * 1024) {
auto kernel = vllm::sampled_topk::sampled_topk_kernel<TopK>;
constexpr size_t smem_size =
vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC;
cudaError_t status = cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
STD_TORCH_CHECK(status == cudaSuccess,
"sampled_topk smem failed: ", cudaGetErrorString(status));
kernel<<<num_rows, vllm::sampled_topk::kThreads, smem_size, stream>>>(
logits.const_data_ptr<float>(), lengths.const_data_ptr<int32_t>(),
output.mutable_data_ptr<int32_t>(), stride,
static_cast<int>(std::min(max_seq_len, logits.size(1))));
} else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
csrc/libtorch_stable/torch_bindings.cpp
+9 -1
View File
@@ -425,12 +425,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"bool is_neox, Tensor position_ids, "
"int forced_token_heads_per_warp=-1) -> ()");
ops.def(
"fused_deepseek_v4_kv_rope_insert("
"Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, "
"Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, "
"bool kv_mxfp8=False) -> "
"()");
ops.def(
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
"Tensor q_in, Tensor kv, Tensor! k_cache, "
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
"int q_head_padded, float eps, int cache_block_size, "
"bool apply_q_norm=True) -> Tensor");
"bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor");
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
@@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
// Positional encoding kernels (shared CUDA/ROCm)
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
ops.impl("fused_deepseek_v4_kv_rope_insert",
TORCH_BOX(&fused_deepseek_v4_kv_rope_insert));
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
ops.impl(
csrc/rocm/skinny_gemms_int4.cu
+41 -18
View File
@@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
return min(a, b);
}
// Reads row `row` of packed zero points [M/8, num_groups] uint32.
//
// The parameters must be signed and the indexing must use shift/mask: with
// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8)
// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239.
__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed,
const int row, const int group,
const int num_groups) {
const uint32_t w = zp_packed[(row >> 3) * num_groups + group];
return (w >> (4 * (row & 7))) & 0xFu;
}
// W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations
// Targets the "sml" case where activations fit in LDS.
// A_CHUNK: number of K-elements processed per thread per step.
@@ -116,7 +128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By,
const uint8_t* B_packed,
const scalar_t* __restrict__ A, const scalar_t* scale,
const scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
uint32_t group_idx = k_ / GROUP_SIZE;
scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
zp_nibble(zero_points, m + y, group_idx, num_groups)));
#pragma unroll
for (uint32_t b = 0; b < A_CHUNK; b++) {
cvtB.h[b] = cvtB.h[b] - zp;
@@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
if constexpr (HAS_ZERO_POINTS) {
uint32_t group_idx_zp = k_ / GROUP_SIZE;
float zp_f = __s2float(
zero_points[(m + y) * num_groups + group_idx_zp]);
float zp_f = static_cast<float>(
zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
partial -= (128.0f + zp_f) * act_sum;
} else {
partial -= 136.0f * act_sum;
@@ -335,7 +348,7 @@ __global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx,
const int By, const uint8_t* B_packed,
const scalar_t* __restrict__ A,
const scalar_t* scale,
const scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
@@ -352,7 +365,7 @@ template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
__global__ void __launch_bounds__(WvPrGrp* THRDS)
wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By,
const uint8_t* B_packed, const scalar_t* __restrict__ A,
const scalar_t* scale, const scalar_t* zero_points,
const scalar_t* scale, const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -508,7 +521,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {
if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
uint32_t group_idx = k_ / GROUP_SIZE;
scalar_t zp = zero_points[(m + y) * num_groups + group_idx];
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
zp_nibble(zero_points, m + y, group_idx, num_groups)));
#pragma unroll
for (uint32_t b = 0; b < A_CHUNK; b++) {
cvtB.h[b] = cvtB.h[b] - zp;
@@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
if constexpr (HAS_ZERO_POINTS) {
uint32_t group_idx_zp = k_ / GROUP_SIZE;
float zp_f = __s2float(
zero_points[(m + y) * num_groups + group_idx_zp]);
float zp_f = static_cast<float>(
zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
partial -= (128.0f + zp_f) * act_sum;
} else {
partial -= 136.0f * act_sum;
@@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx,
const int By, const uint8_t* B_packed,
const scalar_t* __restrict__ A,
const scalar_t* scale,
const scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
@@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) {
// in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32)
// in_b: activations [N, K] (fp16/bf16)
// in_scale: group scales [M, K/group_size] (fp16/bf16)
// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16)
// in_zero_points: optional raw zero points [M/8, K/group_size] (int32,
// 8 uint4 per word along dim 0)
// If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric).
// If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8).
// group_size: 32, 64, or 128
@@ -664,14 +679,22 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
"Scale must be [M, K/group_size] = [", M_in, ", ", num_groups,
"] but got [", in_scale.size(0), ", ", in_scale.size(1), "]");
if (in_zero_points.has_value()) {
TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(),
"Zero points dtype must match activation dtype");
// The kernel reads the words as uint32, so either signedness is accepted.
TORCH_CHECK(in_zero_points->dtype() == at::kInt ||
in_zero_points->dtype() == at::kUInt32,
"Zero points must be int32 or uint32 (packed 8x uint4 along "
"dim 0), got ",
in_zero_points->dtype());
TORCH_CHECK(in_zero_points->dim() == 2,
"Zero points must be 2D [M, K/group_size], got shape ",
"Zero points must be 2D [M/8, K/group_size], got shape ",
in_zero_points->sizes());
TORCH_CHECK(in_zero_points->size(0) == M_in &&
TORCH_CHECK(M_in % 8 == 0,
"M must be divisible by 8 for packed zero points, got ", M_in);
TORCH_CHECK(in_zero_points->is_contiguous(),
"Zero points must be contiguous");
TORCH_CHECK(in_zero_points->size(0) == M_in / 8 &&
in_zero_points->size(1) == num_groups,
"Zero points must be [M, K/group_size] = [", M_in, ", ",
"Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ",
num_groups, "] but got [", in_zero_points->size(0), ", ",
in_zero_points->size(1), "]");
}
@@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
const fptype* aptr = reinterpret_cast<const fptype*>(in_b.data_ptr());
const fptype* sptr =
reinterpret_cast<const fptype*>(in_scale.data_ptr());
const fptype* zpptr =
const uint32_t* zpptr =
in_zero_points.has_value()
? reinterpret_cast<const fptype*>(in_zero_points->data_ptr())
? reinterpret_cast<const uint32_t*>(in_zero_points->data_ptr())
: nullptr;
const fptype* biasptr =
(in_bias.has_value() && in_bias->numel() > 0)
csrc/rocm/torch_bindings.cpp
+2 -1
View File
@@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
// W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,
// optional zero points for asymmetric quantization
// optional zero points [M/8, K/group_size] int32 for asymmetric
// quantization
rocm_ops.def(
"wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "
"Tensor? in_zero_points, Tensor? in_bias, int CuCount, "
docker/Dockerfile
+2 -1
View File
@@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \
git \
curl \
sudo \
elfutils-devel \
rdma-core-devel \
&& dnf clean all \
&& rm -rf /var/cache/dnf
@@ -640,7 +641,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
python3 use_existing_torch.py --prefix; \
fi \
&& if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \
sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \
sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \
requirements/cuda.txt; \
fi
docker/Dockerfile.rock
+9
View File
@@ -875,6 +875,15 @@ RUN mkdir src && mv vllm src/vllm
# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests.
ENV VLLM_GPU_SYNC_CHECK=error
FROM test AS test_smoke
ARG ROCM_SMOKE_ID
RUN --network=none \
bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \
&& printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok
FROM scratch AS export_test_smoke
COPY --from=test_smoke /vllm-smoke-ok /
# -----------------------
# Final vLLM image
FROM mori_base AS final_common
docker/Dockerfile.rock_base
+9
View File
@@ -146,6 +146,15 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \
fi; \
done \
&& for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \
librocprofiler-sdk-roctx.so.1; do \
if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \
&& [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \
rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \
ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \
"${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \
fi; \
done \
&& ldconfig
# Install sccache if USE_SCCACHE is enabled (for release builds)
docs/contributing/README.md
+28
View File
@@ -315,6 +315,34 @@ review process:
the `ready` label, the PR author can use `/ci run`, `/ci retry`, `/ci cancel`,
or the corresponding `/amd-ci` variants. New commits do not start upstream
CI automatically.
- Before creating a Buildkite build, `/ci run` and `/amd-ci run`, including
`all` and `nightly`, check the requested PR commit against its target branch
in the upstream repository. Your branch must contain every commit currently
on that branch: **zero commits behind**. This also applies to PRs targeting
release branches or another PR's branch.
- If the PR is behind, the bot reports the count without creating a new build.
Merge or rebase onto the latest target branch, then rerun the command. If the
target branch advances in the meantime, update your branch and try again.
- To run CI on an outdated branch at your own risk, append `--allow-stale` to
any run command, for example `/ci run --allow-stale` or
`/amd-ci run all --allow-stale`. The same authorization requirements apply.
The bot reports the lag and warns that outdated CI configuration may cause
failures. Before merging, merge or rebase onto the latest target branch and
rerun CI without `--allow-stale` on the latest PR commit.
- Retrying or cancelling builds does not check branch freshness.
- These commands do not modify your branch or enforce merge requirements.
Direct Buildkite launches are outside this workflow. No additional token is
needed.
Maintainers must configure required Buildkite status checks on the latest PR
commit and enable GitHub's
[Require branches to be up to date before merging](https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/managing-protected-branches/about-protected-branches#require-status-checks-before-merging)
protection on the relevant target branches to enforce the merge policy. Select
statuses that confirm the required tests completed, using Buildkite as their
expected source. An aggregate status that reports "passed and blocked" does not
by itself guarantee those tests ran; verify the pipeline's required and optional
blocked steps before choosing statuses or changing blocked-build reporting.
The CI command workflow does not configure these repository or pipeline settings.
### Pull Request Limits and Escalation
docs/contributing/profiling.md
+31 -13
View File
@@ -90,18 +90,18 @@ $ curl -X POST http://localhost:8000/stop_profile
is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a
Chrome trace and works through the same vLLM profiling controls as the PyTorch
and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and
requires eager execution.
supports CUDA graph attribution.
Start a server with a local output directory:
Start a server with a local output directory and graph attribution:
```bash
vllm serve meta-llama/Llama-3.1-8B-Instruct \
--enforce-eager \
--profiler-config '{
"profiler": "proton",
"proton_profiler_dir": "./proton_profile",
"proton_output_format": "hatchet",
"proton_hook": "triton"
"proton_hook": "triton",
"proton_graph_attribution": true
}'
```
@@ -111,7 +111,7 @@ rank-qualified output name, such as
`proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`,
so distributed workers, restarted servers, and repeated profiling runs do not
overwrite one another. A `profile_prefix` is included when supplied. Each
profile is finalized by `/stop_profile` and is ready to inspect immediately.
profile is written by `/stop_profile` and is ready to inspect immediately.
The Proton-specific options are:
@@ -121,21 +121,39 @@ The Proton-specific options are:
- `proton_mode`: an optional backend mode string
- `proton_hook`: `triton` to record Triton launch metadata, or unset
- `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset
- `proton_graph_attribution`: observe CUDA graph capture for replay attribution;
disabled by default and requires `proton_data: "tree"`
`hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while
`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet
supported, so `--enforce-eager` is required for every Proton mode.
`chrome_trace` requires `proton_data: "trace"`.
Automatic backend selection is recommended. vLLM currently supports Proton's
`cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not
expose Proton's experimental instrumentation backend because current upstream
Triton builds can produce profiles without timing metrics. Backend-specific
modes, including `pcsampling`, can be selected with `proton_mode`.
can produce profiles without timing metrics. When `proton_graph_attribution` is
enabled, Proton observes vLLM's CUDA graph capture with the configured profiling
session active, then deactivates that same session until profiling starts. This
lets later profiles attribute replayed kernels without retaining model-startup
activity. Backend-specific modes can be selected with `proton_mode`;
`pcsampling` synchronizes the CUDA context and therefore requires
`--enforce-eager`. When CUDA graphs are enabled (including encoder graphs),
Proton requires `proton_graph_attribution: true` to collect replayed kernels.
For Chrome traces, disable CUDA graphs with `--enforce-eager`.
Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The
`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or
newer; vLLM rejects these options with the detected version before starting a
session.
CUDA graph-attributed profiles support repeated `start_profile`/`stop_profile`
runs. Each stop flushes and writes one tree-data phase while preserving the
graph-aware session. Without `proton_graph_attribution`, each `stop_profile`
instead finalizes and writes an independent Proton session.
CUDA graph attribution requires Triton 3.7 or newer. It uses the phase data API
to discard graph-capture activity and separate profiling runs. The `hatchet_msgpack`
output format and `periodic_flushing` mode also require Triton 3.7 or newer.
`periodic_flushing` cannot be combined with graph attribution because both
manage the session's data phases. Ordinary Proton profiling remains available
with Triton 3.6.
Graph attribution retains capture metadata for the worker lifetime. With eager
execution or no graphs to capture, profiling uses ordinary independent sessions.
Inspect tree profiles with:
docs/contributing/vulnerability_management.md
-2
View File
@@ -33,9 +33,7 @@ on GitHub. However, if you need to contact the VMT directly for an urgent issue,
you may contact the following individuals:
- Simon Mo - <simon.mo@hey.com>
- Russell Bryant - <rbryant@redhat.com>
- Juan Pérez de Algaba - <jperezde@redhat.com>
- Huzaifa Sidhpurwala - <huzaifas@redhat.com>
## Slack Discussion
docs/design/metrics.md
+19 -4
View File
@@ -33,7 +33,8 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met
- `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts.
- `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts.
- `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT).
- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency.
- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency (time between consecutive streamed outputs).
- `vllm:request_time_per_output_token_seconds` (Histogram) - Per-request Time Per Output Token (TPOT).
- `vllm:e2e_request_latency_seconds` (Histogram) - End-to-end request latency.
- `vllm:request_prefill_time_seconds` (Histogram) - Request prefill time.
- `vllm:request_decode_time_seconds` (Histogram) - Request decode time.
@@ -49,7 +50,8 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of
- `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds.
- `vllm:prompt_tokens` - Prompt tokens.
- `vllm:generation_tokens` - Generation tokens.
- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds.
- `vllm:inter_token_latency_seconds` - Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains exactly one token, but differs from request-level TPOT when an output contains multiple tokens or when aggregation weights differ.
- `vllm:request_time_per_output_token_seconds` - Per-request Time Per Output Token (TPOT) in seconds, computed as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. It is recorded as zero for requests that generate no more than one token.
- `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds.
- `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states.
- `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM.
@@ -63,6 +65,20 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of
See [the PR which added this Dashboard](https://github.com/vllm-project/vllm/pull/2316) for interesting and useful background on the choices made here.
`vllm:inter_token_latency_seconds` records one sample per streamed output event
(the wall-clock gap between successive outputs), whereas
`vllm:request_time_per_output_token_seconds` is recorded once per finished request
as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. These two
metrics differ when an output bundles multiple tokens (e.g. speculative
decoding) or when mean ITL and mean per-request TPOT are aggregated over
differing weights. See [Understanding the latency metrics](../benchmarking/cli.md#understanding-the-latency-metrics)
for the benchmark definitions. Requests that generate no more than one token
are recorded with a TPOT of zero, while `vllm bench serve` excludes them from
its TPOT statistics. Therefore, the Prometheus histogram and benchmark TPOT
statistics can differ. Use `vllm:request_time_per_output_token_seconds` for
request-level TPOT and `vllm:inter_token_latency_seconds` when you specifically
want the inter-output latency distribution.
### Prometheus Client Library
Prometheus support was initially added [using the aioprometheus library](https://github.com/vllm-project/vllm/pull/1890), but a switch was made quickly to [prometheus_client](https://github.com/vllm-project/vllm/pull/2730). The rationale is discussed in both linked PRs.
@@ -238,8 +254,7 @@ statistics relating to that iteration:
iteration.
- The prefill intervals for any requests that completed prefill in
this iteration.
- The inter-token intervals (Time Per Output Token, TPOT), for all
requests included in this iteration.
- The inter-token intervals, for all requests included in this iteration.
- The Time-To-First-Token (TTFT) for any requests that completed
prefill in this iteration. However, we calculate this interval
relative to when the request was first received by the frontend
docs/design/nixl_kv_push_connector.md
+44
View File
@@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing
is D's own request id, taken from the registration), so the D-side
accounting code is unchanged.
## Pipeline parallelism and hybrid KV caches
The pull connector requires the local and remote workers to expose a
congruent list of KV regions — region *i* here corresponds to region
*i* there. That assumption breaks under pipeline parallelism (PP)
combined with a hybrid (HMA) KV layout:
* a PP-sharded prefiller (**P**) holds only a slice of the model's
layers while the `PP=1` decoder (**D**) holds them all, so region
counts differ;
* with HMA, several layer names are pooled into one region and the layer
that represents a pooled region can differ between P and D.
`NixlPushConnector` handles this for PP-sharded producers by routing
**by layer-name (member) identity** instead of by region index. Each
worker advertises which layer names back each of its NIXL regions in the
handshake metadata (`NixlAgentMetadata.region_members`). A producer that
needs member routing derives its member-major layout once, when it
registers its KV caches, and every transfer it issues uses that order.
`add_remote_agent` then selects exactly the remote regions this stage
owns and reorders them to match, so both sides stay paired regardless of
how each remote rank happens to order its metadata.
Addresses, block lengths, strides, and per-region capacities follow the same
member order. Descriptor offsets use each member's region capacity, so P and
D need not allocate the same number of blocks. Physical allocations are still
registered once, even when multiple members share them.
Invariants enforced when the remote regions are aligned:
* every locally owned member must be advertised exactly once by the
remote; a missing member fails the handshake rather than silently
leaving that layer's KV stale, and remote-only members (owned by other
PP stages) are ignored;
* a remote that omits member metadata while the local layout requires
member routing fails the handshake instead of falling back to
region-index routing;
* member order is a property of the local layout alone, so the same local
source descriptors serve every remote engine and TP rank; only the
remote descriptor list is rebuilt per rank.
Decode-side PP is unsupported because completions are counted per
consumer rank. Mamba/SSM hybrids are unsupported under PP.
## Scheduler-side responsibilities
`NixlPushConnectorScheduler` extends the base scheduler with:
docs/features/README.md
+1 -1
View File
@@ -64,7 +64,7 @@ th:not(:first-child) {
| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
| [pooling](../models/pooling_models/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
docs/features/kv_offloading_usage.md
+51
View File
@@ -22,6 +22,35 @@ flowchart LR
CPU <--> SN["..."]
```
## Per-request load control
Individual requests can cap how many tokens are loaded from offloaded storage
by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens
beyond those already available in the GPU prefix cache. Set it to `0` to
disable external loading for the request:
```json
{
"kv_transfer_params": {
"max_load_tokens": 0
}
}
```
GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap
are recomputed. Omitting the field leaves loading uncapped. The value must be a
non-negative integer; invalid values are ignored. Positive caps are rounded
down to a boundary supported by the configured KV cache groups. Offloading
remains enabled unless it is controlled separately with `max_offload_tokens`.
!!! warning
`max_load_tokens` is experimental and subject to change.
`kv_load_tiers` continues to select secondary tiers. CPU is always included:
it can satisfy a resident hit directly and is the required staging tier for
secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but
does not query any secondary tier.
## Terminology: Chunks
The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers.
@@ -203,9 +232,31 @@ Block content hashes must match across instances for peers to exchange blocks (s
| `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. |
| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. |
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.
A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives
within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier
slots. Raise it when prefills legitimately take longer than the default:
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "TieringOffloadingSpec",
"secondary_tiers": [
{
"type": "p2p",
"unbound_store_timeout_s": 180
}
]
}
}'
```
#### 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.
docs/features/nixl_connector_compatibility.md
+18
View File
@@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
- `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting.
- Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
### Pipeline parallelism
The default (pull) NixlConnector does **not** support
`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA)
layout: region indices are not a stable identity across the
prefill/decode layer split, so the connector raises at startup. Use the
push connector (`NixlPushConnector`) for pipeline parallelism with
hybrid KV caches — it routes transfers by layer-name (member) identity,
so a PP-sharded prefiller can write into a `PP=1` decoder. See
[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md).
Current push PP + HMA limitations:
- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported.
- Hybrid SSM/Mamba layouts are not supported under PP.
- HMA requires the same block size on P and D.
- Attention-HMA member routing requires decode TP to be no greater than prefill TP.
### Quantized KV cache
[Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake.
docs/features/quantization/README.md
+17
View File
@@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM:
- [TorchAO](torchao.md)
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
## Selecting Linear Backends per Quantization
`--linear-backend` selects one backend for all quantized linear layers. For
mixed-precision models that use more than one linear quantization scheme, use
`linear_backend_per_quant` to override the backend for individual schemes:
```bash
vllm serve <model> \
--linear-backend cutlass \
--kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}'
```
Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear
layers use CUTLASS. Per-quantization overrides take precedence over
`--linear-backend`; schemes without an override continue to use the global
setting, including automatic selection when it is `auto`.
## Supported Hardware
The table below shows the compatibility of various quantization implementations with different hardware platforms in vLLM:
docs/features/watermarking.md
+6 -1
View File
@@ -29,7 +29,12 @@ from vllm import SamplingParams
sampling_params = SamplingParams(watermarking=False)
```
The OpenAI-compatible APIs accept the same `watermarking: false` request field.
The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the
same `watermarking: false` request field. The Rust token API accepts it in
`sampling_params`, and the Rust gRPC `GenerateRequest` accepts
the optional `watermarking` field. Omitting the field keeps watermarking enabled
when the engine has a watermark configuration.
Deployments that require watermarking must restrict this field to trusted
callers, or strip and validate it at the ingress boundary, so untrusted clients
cannot opt out.
docs/governance/committers.md
+5 -4
View File
@@ -32,6 +32,7 @@ Sorted alphabetically by GitHub handle:
- [@ivanium](https://github.com/ivanium): KV Connector and offload
- [@jeejeelee](https://github.com/jeejeelee): LoRA, new model support
- [@jikunshang](https://github.com/jikunshang): Intel CPU/XPU integration
- [@jperezdealgaba](https://github.com/jperezdealgaba): Security
- [@khluu](https://github.com/khluu): CI infrastructure
- [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
- [@markmc](https://github.com/markmc): Observability
@@ -44,7 +45,6 @@ Sorted alphabetically by GitHub handle:
- [@pavanimajety](https://github.com/pavanimajety): NVIDIA GPU integration
- [@ProExpertProg](https://github.com/ProExpertProg): Compilation, startup UX
- [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
- [@russellb](https://github.com/russellb): Structured output, engine core, security
- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
- [@shen-shanshan](https://github.com/shen-shanshan): AMD GPU / ROCm integration, multimodality, ViT CUDA graph
- [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
@@ -82,6 +82,7 @@ Committers who have contributed to vLLM significantly in the past (thank you!) b
- [@pcmoritz](https://github.com/pcmoritz): MoE
- [@rkooo567](https://github.com/rkooo567): Chunked prefill
- [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
- [@russellb](https://github.com/russellb): Structured output, engine core, security
- [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
- [@sroy745](https://github.com/sroy745): Speculative decoding
- [@Yard1](https://github.com/Yard1): kernels and performance
@@ -100,7 +101,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- KV Cache Manager: memory management layer within scheduler maintaining KV cache logical block data
- @heheda12345, @WoosukKwon
- AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
- @robertgshaw2-redhat, @njhill, @russellb
- @robertgshaw2-redhat, @njhill
- ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
- KV Connector: Connector interface and implementation for KV cache offload and transfer
@@ -153,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
- @WoosukKwon, @benchislett, @MatthewBonanni
- Structured Output: The structured output implementation
- @russellb, @aarnphm
- @aarnphm
- RL: The RL related features such as collective rpc, sleep mode, etc.
- @youkaichao, @zhuohan123
- LoRA: @jeejeelee
@@ -167,7 +168,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- Documentation: @hmellor, @DarkLight1337, @simon-mo
- Benchmarks: @ywang96, @simon-mo
- CI, Build, Release Process: @khluu, @njhill, @simon-mo, @Harry-Chen, @vadiklyutiy
- Security: @russellb
- Security: @jperezdealgaba
### External Kernels Integration
docs/models/supported_models.md
+8
View File
@@ -490,6 +490,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
| `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
@@ -601,6 +602,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
| `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I<sup>+</sup> + V<sup>+</sup> | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ |
| `NemotronH_Nano_Omni_Reasoning_V3`, `NemotronH_Nano_VL_V2` | Nemotron Nano 3 Omni, Nemotron Nano V2 VL | T + I<sup>E+</sup> + V<sup>+</sup> + A<sup>*</sup> | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | ✅︎ | |
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
| `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
@@ -700,6 +702,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor
`meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through
its [DFlash](../features/speculative_decoding/README.md) path.
!!! note
Both Nemotron Nano architecture names map to the same vLLM implementation.
Audio inputs are only supported by the Omni variants; `NemotronH_Nano_VL_V2`
checkpoints such as `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` accept text,
image and video only.
!!! note
The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>
docs/serving/online_serving/README.md
+7 -7
View File
@@ -121,12 +121,7 @@ For further details on profiling vLLM, please refer to [this page](../../contrib
## Scale-Out APIs
Scale-out APIs are disabled by default on `vllm serve`. The environment
variable accepts only `0` or `1`; set `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1` to
register the endpoints below. The
dedicated `vllm launch render` and `vllm serve --tokens-only` modes are explicit
opt-ins and enable their required endpoints when the variable is unset; an
explicit value of `0` is rejected for those modes.
Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` to register the endpoints below. The dedicated `vllm launch render` and `vllm serve --tokens-only` modes always register their required endpoints regardless of `--enable-scale-out`.
### Tokens IN <> Tokens OUT APIs
@@ -135,7 +130,12 @@ explicit value of `0` is rejected for those modes.
### Renderer APIs
For further details on renderer APIs, please refer to [this page](renderer.md).
Renderer APIs preprocess completion, chat, and Responses requests without running
inference. They handle tokenization, model-specific prompt formatting, and
multimodal preprocessing, returning prompt token IDs, sampling parameters, and
any processed multimodal inputs for generation.
See the [renderer guide](renderer.md) for setup instructions and examples.
- [Completions Render API](renderer.md) (`/v1/completions/render`)
- Render completion requests
docs/serving/online_serving/derenderer.md
+1 -2
View File
@@ -60,8 +60,7 @@ Launch the two servers first:
```bash
vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve \
meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200
vllm serve meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200
```
```python
docs/serving/online_serving/openai_compatible_server.md
+5
View File
@@ -199,6 +199,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../../examples/tool_calling/openai_responses_client_with_tools.py)
To get prompt token IDs before generation, use
[`/v1/responses/render`](renderer.md#get-responses-prompt-token-ids). It preprocesses
a self-contained Responses request without inference and returns the prompt
token IDs, so a separate `/tokenize` call is not needed.
#### Extra parameters
The following extra parameters in the request object are supported:
docs/serving/online_serving/renderer.md
+35 -6
View File
@@ -7,16 +7,14 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and
- Tokens-in / tokens-out engine: Make the engine a pure token-in / token-out service, decoupled from request preprocessing.
The dedicated `vllm launch render` server always exposes the `/render` and
`/derender` endpoints when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` is unset or set to
`1`. An explicit value of `0` conflicts with the renderer command and is
rejected at startup.
`/derender` endpoints.
Scale-out endpoints, including `/render`, `/derender`, and
`/inference/v1/generate`, are disabled by default on a standard inference
server. To expose them with `vllm serve`, opt in explicitly:
```bash
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve <model>
vllm serve <model> --enable-scale-out
```
## API Reference
@@ -28,27 +26,58 @@ VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve <model>
- [Responses Render API](renderer.md) (`/v1/responses/render`)
- Render a self-contained Responses request
## Get Responses prompt token IDs
Use `/v1/responses/render` to get prompt token IDs before choosing a model
replica. Rendering applies prompt construction and preprocessing without running
inference.
The Responses render endpoint uses the same prompt construction as
`/v1/responses` and returns one token-in `GenerateRequest`. It is stateless:
inline history is supported, but `previous_response_id` is not. Callers must
resolve stored response state and include the resulting history in the request
before rendering.
Configure the renderer and generation workers with the same model, tokenizer,
chat template, and preprocessing options. Render the full request, including
instructions, history, tools, and any template or truncation options, so the
returned IDs reflect the prompt the model will receive.
For multimodal requests, the `GenerateRequest` contains the model-processed
multimodal payload, which can be substantially larger than the source image or
video. The caller must forward that payload unchanged to the generation service
and provision transport limits and memory accordingly.
For example, start a standard inference server with scale-out endpoints enabled:
```bash
curl http://localhost:8000/v1/responses/render \
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
vllm serve meta-llama/Llama-3.1-8B-Instruct
```
Send a Responses request and use `jq` to extract the `token_ids` field:
```bash
curl --fail --silent --show-error http://localhost:8000/v1/responses/render \
-H "Content-Type: application/json" \
-d '{
"model": "meta-llama/Llama-3.1-8B-Instruct",
"input": "Explain prefix caching in one sentence.",
"max_output_tokens": 32
}'
}' | jq '.token_ids'
```
To count the rendered prompt tokens for this text request, replace the `jq`
filter with `'.token_ids | length'`. The endpoint returns the full
`GenerateRequest`; `jq` filters the response on the client. Keep the complete
response if you will forward it to `/inference/v1/generate`, including any
multimodal features.
If the server has `--api-key` or `VLLM_API_KEY` configured, add
`-H "Authorization: Bearer <api-key>"` to the request. See
[API key authentication limitations](../../usage/security.md#api-key-authentication-limitations)
for the existing authentication boundaries.
For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md).
## Multimodal Render Features
docs/usage/security.md
+10 -10
View File
@@ -158,27 +158,27 @@ When `--api-key` is configured, the following endpoints require Bearer token aut
- `/v1/models` - List available models
- `/v1/chat/completions` - Chat completions
- `/v1/chat/completions/batch` - Batch chat completions
- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
- `/v1/completions` - Text completions
- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
- `/v1/embeddings` - Generate embeddings
- `/v1/audio/transcriptions` - Audio transcription
- `/v1/audio/translations` - Audio translation
- `/v1/messages` - Anthropic-compatible messages API
- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`)
- `/v1/messages/count_tokens` - Count tokens for Anthropic messages
- `/v1/responses` - Create a response
- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled)
- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render` unless explicitly disabled)
- `/v1/responses/{response_id}` - Retrieve a response
- `/v1/responses/{response_id}/cancel` - Cancel a response
- `/v1/score` - Scoring API
- `/v1/rerank` - Reranking API
- `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
- `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
- `/inference/v1/generate` - Generate completions (available when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or with `--tokens-only` unless explicitly disabled)
- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires both `VLLM_ENABLE_COHERE_API=1` and `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`)
- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`)
- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`)
- `/v2/embed` - Cohere Embed API
- `/v2/rerank` - Cohere Rerank API
@@ -206,11 +206,11 @@ The following endpoints **do not require authentication** even when `--api-key`
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
- `/update_weights` - Update model weights (can alter model behavior)
- `/get_world_size` - Get distributed world size
- `/abort_requests` - Abort in-flight requests (available with `--tokens-only` unless `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0`)
- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`)
**Utility endpoints:**
- `/tokenize` - Tokenize text (not disabled by `VLLM_ENABLE_SCALE_OUT_ENDPOINTS`)
- `/tokenize` - Tokenize text (not gated by `--enable-scale-out`)
- `/detokenize` - Detokenize tokens
- `/health` - Health check
- `/ping` - SageMaker health check
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+9 -2
View File
@@ -278,7 +278,6 @@ async def fanout_encoder_primer(
item_transfer_ids[idx] = transfer_id
encoder_req = {
# You *may* need to keep additional fields
"model": orig_request.get("model"),
"messages": [
{
@@ -292,6 +291,14 @@ async def fanout_encoder_primer(
# once the prompt is encoded and its embeddings are published.
"stream": False,
}
for key in (
"mm_processor_kwargs",
"media_io_kwargs",
"priority",
"session_id",
):
if key in orig_request:
encoder_req[key] = orig_request[key]
if consumer_zmq is not None:
# No mm_hash here on purpose. The encoder's own
# `mm_features[i].identifier` is derived from the uuid *and* the
@@ -378,7 +385,7 @@ async def fanout_encoder_primer(
# connector's own handle on the published embedding (for NIXL,
# peer_host/peer_port/size_bytes). The decoder's connector
# looks it up by mm_hash on the request, so carry it through.
ec_params[item_uuids.get(idx, ec_mm_hash)] = reported
ec_params[ec_mm_hash] = reported
if NO_REWRITE and consumer_zmq is not None:
ec_params.setdefault("ec_items", []).append(
{"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]}
examples/observability/prometheus_grafana/grafana.json
+167 -1
View File
@@ -318,7 +318,7 @@
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"description": "Inter token latency in seconds.",
"description": "Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains one token, but differs from request-level TPOT under speculative decoding or varying aggregation weights.",
"fieldConfig": {
"defaults": {
"color": {
@@ -479,6 +479,172 @@
"title": "Inter Token Latency",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"description": "Per-request Time Per Output Token (TPOT) in seconds, computed as (end-to-end latency - TTFT) divided by the number of output tokens minus one. Requests with no more than one output token are recorded as zero; vllm bench serve excludes these requests from its TPOT statistics.",
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"axisBorderShow": false,
"axisCenteredZero": false,
"axisColorMode": "text",
"axisLabel": "",
"axisPlacement": "auto",
"barAlignment": 0,
"barWidthFactor": 0.6,
"drawStyle": "line",
"fillOpacity": 0,
"gradientMode": "none",
"hideFrom": {
"legend": false,
"tooltip": false,
"viz": false
},
"insertNulls": false,
"lineInterpolation": "linear",
"lineWidth": 1,
"pointSize": 5,
"scaleDistribution": {
"type": "linear"
},
"showPoints": "auto",
"spanNulls": false,
"stacking": {
"group": "A",
"mode": "none"
},
"thresholdsStyle": {
"mode": "off"
}
},
"mappings": [],
"thresholds": {
"mode": "absolute",
"steps": [
{
"color": "green",
"value": null
},
{
"color": "red",
"value": 80
}
]
},
"unit": "s"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 48
},
"id": 17,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "single",
"sort": "none"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"disableTextWrap": false,
"editorMode": "builder",
"expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
"fullMetaSearch": false,
"includeNullMetadata": false,
"instant": false,
"legendFormat": "P99",
"range": true,
"refId": "A",
"useBackend": false
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"disableTextWrap": false,
"editorMode": "builder",
"expr": "histogram_quantile(0.95, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
"fullMetaSearch": false,
"hide": false,
"includeNullMetadata": false,
"instant": false,
"legendFormat": "P95",
"range": true,
"refId": "B",
"useBackend": false
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"disableTextWrap": false,
"editorMode": "builder",
"expr": "histogram_quantile(0.9, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
"fullMetaSearch": false,
"hide": false,
"includeNullMetadata": false,
"instant": false,
"legendFormat": "P90",
"range": true,
"refId": "C",
"useBackend": false
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"disableTextWrap": false,
"editorMode": "builder",
"expr": "histogram_quantile(0.5, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))",
"fullMetaSearch": false,
"hide": false,
"includeNullMetadata": false,
"instant": false,
"legendFormat": "P50",
"range": true,
"refId": "D",
"useBackend": false
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "rate(vllm:request_time_per_output_token_seconds_sum{model_name=\"$model_name\"}[$__rate_interval])\n/\nrate(vllm:request_time_per_output_token_seconds_count{model_name=\"$model_name\"}[$__rate_interval])",
"hide": false,
"instant": false,
"legendFormat": "Mean",
"range": true,
"refId": "E"
}
],
"title": "Time Per Output Token (TPOT)",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
examples/scale_out/example_mm_serve.py
+3 -3
View File
@@ -13,9 +13,9 @@ zero client-side transformation.
Launch the server first:
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
vllm serve Qwen/Qwen3-VL-2B-Instruct \
--dtype bfloat16 --max-model-len 4096 --enforce-eager
vllm serve Qwen/Qwen3-VL-2B-Instruct \
--dtype bfloat16 --max-model-len 4096 --enforce-eager \
--enable-scale-out
Then run this script:
examples/scale_out/token_generation_client.py
+1 -1
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Call a token-in/token-out server started with:
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve Qwen/Qwen3-0.6B
vllm serve Qwen/Qwen3-0.6B --enable-scale-out
"""
import httpx
requirements/rubin-prerelease.txt
+16 -6
View File
@@ -2,15 +2,25 @@
# Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile).
# TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels
# (vllm-project/vllm#49735).
--pre
--extra-index-url https://download.pytorch.org/whl/nightly/cu134
--extra-index-url https://flashinfer.ai/whl/cu134
--extra-index-url https://flashinfer.ai/whl
torch==2.15.0.dev20260903+cu134
torchvision==0.30.0.dev20260903+cu134
torchaudio==2.11.0.dev20260903+cu134
# Match each architecture's torch pin to the latest torchvision wheel.
torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64"
torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64"
torchvision==0.30.0.dev20260909+cu134
torchaudio==2.11.0.dev20260911+cu134
# docker/Dockerfile strips these two from requirements/cuda.txt whenever
# docker/Dockerfile strips these packages from requirements/cuda.txt whenever
# INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to
# still get installed.
nvidia-cutlass-dsl[cu13]==4.8.0.dev0
flashinfer-python==0.6.18
flashinfer-python==0.6.18.post1
quack-kernels==0.6.5
cuda-python==13.4.1
cuda-bindings==13.4.1
flashinfer-jit-cache==0.6.18.post1+cu134
flashinfer-cubin==0.6.18.post1
nixl==1.4.1
requirements/test/xpu.txt
+2 -17
View File
@@ -257,7 +257,7 @@ googleapis-common-protos==1.75.1
# google-api-core
# opentelemetry-exporter-otlp-proto-grpc
# opentelemetry-exporter-otlp-proto-http
gpt-oss==0.0.9
gpt-oss==0.0.8
# via -r requirements/test/xpu.in
graphql-core==3.2.11
# via hypothesis-graphql
@@ -406,10 +406,6 @@ jsonschema-rs==0.49.9
# via schemathesis
jsonschema-specifications==2025.9.1
# via jsonschema
jupyter-client==8.10.0
# via gpt-oss
jupyter-core==5.9.1
# via jupyter-client
kiwisolver==1.5.1
# via matplotlib
lark==1.2.2
@@ -672,9 +668,7 @@ pillow==12.3.0
# mistral-common
# torchvision
platformdirs==4.11.3
# via
# jupyter-core
# pooch
# via pooch
pluggy==1.6.0
# via
# pytest
@@ -813,7 +807,6 @@ pytest-timeout==2.4.0
python-dateutil==2.9.0.post0
# via
# botocore
# jupyter-client
# matplotlib
# pandas
# typepy
@@ -849,7 +842,6 @@ pyzmq==27.1.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# jupyter-client
rapidfuzz==3.14.5
# via
# -r requirements/test/xpu.in
@@ -1074,8 +1066,6 @@ torchvision==0.28.0+xpu
# via
# open-clip-torch
# timm
tornado==6.5.8
# via jupyter-client
tqdm==4.70.0
# via
# -r requirements/test/../common.txt
@@ -1091,10 +1081,6 @@ tqdm==4.70.0
# pqdm
# sentence-transformers
# transformers
traitlets==5.16.1
# via
# jupyter-client
# jupyter-core
transformers==5.16.1
# via
# -c requirements/common.txt
@@ -1146,7 +1132,6 @@ typing-extensions==4.16.0
# grpcio
# httpx2
# huggingface-hub
# jupyter-client
# lm-eval
# mcp
# mcp-types
rust/Cargo.lock
Generated
+2
View File
@@ -5813,6 +5813,7 @@ dependencies = [
"vllm-engine-core-client",
"vllm-managed-engine",
"vllm-server",
"vllm-text",
"vllm-tracing",
]
@@ -5972,6 +5973,7 @@ dependencies = [
"libc",
"llm-multimodal",
"openssl",
"prost",
"prost-types",
"rmp-serde",
"rmpv",
rust/proto/README.md
+5 -10
View File
@@ -2,14 +2,10 @@
This directory is the canonical source for vLLM's gRPC schema.
The schema is published to `buf.build/vllm-project/vllm`:
- A daily workflow publishes the latest Git `main` schema to the `nightly` label.
- The workflow can be run manually to retry nightly publication.
- Tags matching `v*` update the Buf `main` label and publish the corresponding release label.
- Buf commits and generated SDK versions are immutable and can be pinned by consumers.
Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically.
Schema updates are no longer published to the Buf Schema Registry. Rust consumers
should use `vllm-proto` from crates.io; consumers in other languages can generate
bindings from the `.proto` files in this directory. Buf still builds and lints
the schemas on pull requests, including on forks.
## Rust crate
@@ -47,8 +43,7 @@ change PR. This check becomes available after the first manual publication.
The `proto-crate.yml` workflow checks that the tag matches the crate version,
verifies the package, and publishes it to crates.io.
Use `proto-v*` tags so they do not trigger Buf's `v*` release workflow. There is
no scheduled crate publication. Pull requests and manual workflow runs only
There is no scheduled crate publication. Pull requests and manual workflow runs only
verify packaging, including on forks. Retry a failed release by rerunning its
tag workflow; crates.io versions cannot be overwritten.
rust/proto/buf.md
-8
View File
@@ -1,8 +0,0 @@
# vLLM gRPC API
Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend.
- `Inference` provides unary and streaming generation.
- `Control` provides discovery and request control.
The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds.
rust/proto/buf.yaml
-1
View File
@@ -1,7 +1,6 @@
version: v2
modules:
- path: .
name: buf.build/vllm-project/vllm
lint:
use:
- MINIMAL
rust/proto/inference.proto
+3
View File
@@ -53,6 +53,9 @@ message GenerateRequest {
// Loaded LoRA name; empty selects the base model.
string lora_name = 15;
// Apply the engine's configured watermark. Defaults to true when omitted.
optional bool watermarking = 16;
}
message RandomSampling {
rust/src/bench/README.md
+1
View File
@@ -667,6 +667,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n
- `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`.
- `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns.
- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`.
**Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0.
rust/src/bench/src/config.rs
+32 -6
View File
@@ -549,13 +549,15 @@ impl BenchConfig {
));
}
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply.
if args.dataset_name == DatasetName::ShareGpt {
// Normalize and validate min/max turns. Conversation datasets only
// consume max_turns (the loader walks all available turns up to the
// cap), so the min/num/max coupling used for synthetic generation
// does not apply.
if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) {
if args.multi_turn_max_turns == 1 {
return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
"--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn"
.into(),
));
}
} else {
@@ -722,7 +724,7 @@ impl BenchConfig {
multi_turn_min_turns,
multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if args.multi_turn
&& args.dataset_name == DatasetName::ShareGpt
&& matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf)
&& args.multi_turn_max_turns != 0
{
Some(args.multi_turn_max_turns)
@@ -1067,6 +1069,30 @@ mod tests {
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
}
#[test]
fn test_hf_multi_turn_uses_conversation_turn_cap() {
let args = vec![
"vllm-bench",
"--backend",
"openai-chat",
"--multi-turn",
"--model",
"test-model",
"--dataset-name",
"hf",
"--dataset-path",
"org/sharegpt-dataset",
"--hf-subset",
"sharegpt",
"--multi-turn-max-turns",
"2",
];
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
}
#[test]
fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args();
rust/src/bench/src/datasets/multi_turn.rs
+57
View File
@@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn(
.as_array()
.ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?;
load_sharegpt_multi_turn_from_rows(
tokenizer,
entries,
num_conversations,
output_len_override,
max_turns,
seed,
request_id_prefix,
)
}
/// Load multi-turn conversations from deserialized ShareGPT rows.
pub fn load_sharegpt_multi_turn_from_rows(
tokenizer: &TokenizerKind,
entries: &[serde_json::Value],
num_conversations: usize,
output_len_override: Option<usize>,
max_turns: Option<usize>,
seed: u64,
request_id_prefix: &str,
) -> Result<Vec<MultiTurnConversation>> {
// Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant)
let mut filtered: Vec<&serde_json::Value> = entries
.iter()
@@ -507,6 +528,42 @@ fn gen_prompt_to_target_len(
mod tests {
use super::*;
#[test]
fn test_load_sharegpt_multi_turn_from_hf_rows() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let rows = vec![serde_json::json!({
"conversations": [
{"from": "human", "value": "hello"},
{"from": "gpt", "value": "hi"},
{"from": "human", "value": "how are you?"},
{"from": "gpt", "value": "great"}
]
})];
let conversations =
load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-")
.unwrap();
let actual = conversations
.iter()
.map(|conversation| {
(
conversation.conversation_id.as_str(),
conversation
.turns
.iter()
.map(|turn| (turn.user_message.as_ref(), turn.expected_output_len))
.collect::<Vec<_>>(),
)
})
.collect::<Vec<_>>();
assert_eq!(
actual,
vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])]
);
}
fn common_prefix_bytes(strings: &[&str]) -> usize {
if strings.is_empty() {
return 0;
rust/src/bench/src/multi_turn.rs
+24 -3
View File
@@ -186,9 +186,30 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
));
}
DatasetName::Hf => {
return Err(BenchError::Config(
"HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
));
let tok = tokenizer
.as_ref()
.ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?;
let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
BenchError::Config("--dataset-path is required for --dataset-name hf".into())
})?;
let (rows, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset(
dataset_id,
config.hf_subset.as_deref(),
config.hf_split.as_deref(),
config.num_prompts,
config.seed,
config.disable_shuffle,
)
.await?;
crate::datasets::multi_turn::load_sharegpt_multi_turn_from_rows(
tok,
&rows,
config.num_prompts,
config.hf_output_len,
config.sharegpt_multi_turn_max_turns,
config.seed,
&config.request_id_prefix,
)?
}
};
rust/src/chat/src/backend/hf.rs
+9 -13
View File
@@ -129,7 +129,8 @@ pub(super) async fn load_model_backends(
model_id: &str,
options: LoadModelBackendsOptions,
) -> Result<LoadedModelBackends> {
let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?;
let mut files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?;
files.apply_overrides(&options.hf_overrides)?;
let text_backend = HfTextBackend::from_resolved_model_files(
files.clone(),
model_id.to_string(),
@@ -166,13 +167,11 @@ fn resolve_multimodal_render_info(
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::tempdir;
use thiserror_ext::AsReport as _;
use vllm_text::Prompt;
use vllm_text::backend::hf::TokenizerSource;
use vllm_text::tokenizer::DynTokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
@@ -206,16 +205,12 @@ mod tests {
write_json(&config_path, config_json);
write_json(&tokenizer_config_path, tokenizer_config_json);
vllm_text::backend::hf::ResolvedModelFiles {
tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")),
tokenizer_config_path: Some(tokenizer_config_path),
generation_config_path: None,
preprocessor_config_path: None,
video_preprocessor_config_path: None,
processor_config_path: None,
chat_template_path: None,
config_path: Some(config_path),
}
write_json(&root.join("tokenizer.json"), "{}");
futures::executor::block_on(vllm_text::backend::hf::ResolvedModelFiles::new(
root.to_str().unwrap(),
None,
))
.unwrap()
}
fn test_tokenizer() -> DynTokenizer {
@@ -232,6 +227,7 @@ mod tests {
"test-model".to_string(),
LoadModelBackendsOptions {
revision: None,
hf_overrides: Default::default(),
generation_config: Default::default(),
renderer,
language_model_only: false,
rust/src/chat/src/backend/mod.rs
+3
View File
@@ -5,6 +5,7 @@ use std::collections::HashMap;
use std::sync::Arc;
use serde_json::Value;
use vllm_text::backend::hf::HfOverrides;
use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend};
use crate::error::Result;
@@ -63,6 +64,8 @@ pub type DynChatTextBackend = Arc<dyn ChatTextBackend>;
pub struct LoadModelBackendsOptions {
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
pub revision: Option<String>,
/// JSON Merge Patch applied to the model config before loading any backend.
pub hf_overrides: HfOverrides,
/// Which generation-config sampling defaults to inherit.
pub generation_config: GenerationConfigMode,
/// Which chat renderer implementation to use.
rust/src/chat/src/lib.rs
+7 -3
View File
@@ -34,8 +34,8 @@ pub use parser::{ParserSelection, validate_parser_overrides};
pub use renderer::hf::ChatTemplateContentFormatOption;
pub use renderer::{
ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer,
DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt,
RendererSelection,
DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource,
RenderedPrompt, RendererSelection,
};
pub use request::{
ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
@@ -112,6 +112,10 @@ impl ChatRequestProcessor {
request: &ChatRequest,
rendered: RenderedPrompt,
) -> Result<(Prompt, Option<MmFeatures>)> {
let media_is_empty = match &rendered.media_order {
Some(media_order) => media_order.is_empty(),
None => !request.has_multimodal(),
};
match self.model_dtype {
Some(model_dtype) => {
multimodal::finalize_rendered_prompt(
@@ -122,7 +126,7 @@ impl ChatRequestProcessor {
)
.await
}
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
None if media_is_empty => Ok((rendered.prompt, None)),
None => Err(Error::UnsupportedMultimodalRenderer),
}
}
rust/src/chat/src/multimodal.rs
+152 -53
View File
@@ -16,7 +16,7 @@ use std::fs;
use std::path::Path;
use std::sync::{Arc, LazyLock};
use itertools::izip;
use itertools::{Either, izip};
use llm_multimodal::{
AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout,
ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata,
@@ -33,7 +33,7 @@ use vllm_text::Prompt;
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use crate::error::{Error, Result, bail_multimodal, multimodal};
use crate::renderer::RenderedPrompt;
use crate::renderer::{MediaPartSource, RenderedPrompt};
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
mod audio;
@@ -549,7 +549,8 @@ pub(crate) async fn finalize_rendered_prompt(
info: Option<&MultimodalModelInfo>,
model_dtype: ModelDtype,
) -> Result<(Prompt, Option<MmFeatures>)> {
if !request.has_multimodal() {
let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?;
if media_parts.is_empty() {
return Ok((rendered.prompt, None));
}
let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?;
@@ -561,63 +562,92 @@ pub(crate) async fn finalize_rendered_prompt(
.map_err(|error| multimodal!("{error}"))?,
Prompt::TokenIds(token_ids) => token_ids,
};
let media_parts = extract_media_parts(request)?;
let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;
Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
}
/// Extract media parts from chat messages in message/content order.
///
/// Assistant history is skipped because generated assistant blocks are already
/// represented as text for prompt rendering in this crate.
fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
let mut all_parts = Vec::new();
for message in &request.messages {
let content = match message {
ChatMessage::System { content }
| ChatMessage::Developer { content, .. }
| ChatMessage::User { content }
| ChatMessage::ToolResponse { content, .. } => content,
ChatMessage::Assistant { .. } => continue,
};
let ChatContent::Parts(parts) = content else {
continue;
};
for part in parts {
match part {
ChatContentPart::Text { .. } => {}
ChatContentPart::ImageUrl {
image_url,
detail,
uuid,
} => all_parts.push(MediaContentPart::ImageUrl {
url: image_url.clone(),
detail: *detail,
uuid: uuid.clone(),
}),
ChatContentPart::VideoUrl { video_url, uuid } => {
all_parts.push(MediaContentPart::VideoUrl {
url: video_url.clone(),
uuid: uuid.clone(),
})
}
ChatContentPart::InputAudio { data, format, uuid } => {
all_parts.push(MediaContentPart::AudioUrl {
url: input_audio_data_url(data, format.as_deref())?,
uuid: uuid.clone(),
})
}
ChatContentPart::AudioUrl { audio_url, uuid } => {
all_parts.push(MediaContentPart::AudioUrl {
url: audio_url.clone(),
uuid: uuid.clone(),
})
/// Resolve media parts in the placeholder order reported by the renderer.
fn extract_media_parts(
request: &ChatRequest,
media_order: Option<&[MediaPartSource]>,
) -> Result<Vec<MediaContentPart>> {
let parts = match media_order {
// Resolve exactly the sources selected by the renderer, in placeholder order.
Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> {
let message = request.messages.get(source.message_index).ok_or_else(|| {
multimodal!(
"renderer reported missing chat message {} for multimodal content",
source.message_index
)
})?;
let content = match message {
ChatMessage::System { content }
| ChatMessage::Developer { content, .. }
| ChatMessage::User { content }
| ChatMessage::ToolResponse { content, .. } => content,
ChatMessage::Assistant { .. } => {
bail_multimodal!("renderer reported multimodal assistant content")
}
};
let ChatContent::Parts(parts) = content else {
bail_multimodal!("renderer reported multimodal content in a text-only chat message")
};
let part = parts.get(source.content_part_index).ok_or_else(|| {
multimodal!(
"renderer reported missing content part {} in chat message {}",
source.content_part_index,
source.message_index
)
})?;
Ok(part)
})),
// Fall back to all media in request message and content-part order (e.g. Jinja).
None => Either::Right(request.messages.iter().flat_map(|message| {
let content = match message {
ChatMessage::System { content }
| ChatMessage::Developer { content, .. }
| ChatMessage::User { content }
| ChatMessage::ToolResponse { content, .. } => Some(content),
ChatMessage::Assistant { .. } => None,
};
match content {
Some(ChatContent::Parts(parts)) => parts.as_slice(),
_ => &[],
}
}
}
Ok(all_parts)
.iter()
.filter(|part| part.is_multimodal())
.map(Ok)
})),
};
parts
.map(|part| match part? {
ChatContentPart::Text { .. } => {
bail_multimodal!("renderer reported a text part as multimodal content")
}
ChatContentPart::ImageUrl {
image_url,
detail,
uuid,
} => Ok(MediaContentPart::ImageUrl {
url: image_url.clone(),
detail: *detail,
uuid: uuid.clone(),
}),
ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl {
url: video_url.clone(),
uuid: uuid.clone(),
}),
ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl {
url: input_audio_data_url(data, format.as_deref())?,
uuid: uuid.clone(),
}),
ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl {
url: audio_url.clone(),
uuid: uuid.clone(),
}),
})
.collect()
}
/// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`.
@@ -1043,6 +1073,75 @@ mod tests {
assert!(input_audio_data_url("AAAA", Some("flac")).is_err());
}
#[test]
fn extract_media_parts_uses_request_order_or_explicit_selection() {
let request = ChatRequest {
messages: vec![
ChatMessage::user(vec![
ChatContentPart::text("before image"),
ChatContentPart::ImageUrl {
image_url: "data:image/png;base64,image-b".to_string(),
detail: None,
uuid: Some("image-b".to_string()),
},
]),
ChatMessage::user(vec![ChatContentPart::ImageUrl {
image_url: "data:image/png;base64,image-a".to_string(),
detail: None,
uuid: Some("image-a".to_string()),
}]),
ChatMessage::assistant_text("answer"),
ChatMessage::user("follow-up"),
],
..ChatRequest::for_test()
};
let media_order = [
MediaPartSource {
message_index: 1,
content_part_index: 0,
},
MediaPartSource {
message_index: 0,
content_part_index: 1,
},
];
let selections = [None, Some(&media_order[..]), Some(&[][..])];
let uuids = selections.map(|order| {
extract_media_parts(&request, order)
.unwrap()
.into_iter()
.map(|part| match part {
MediaContentPart::ImageUrl { uuid, .. } => uuid,
_ => panic!("expected image URL"),
})
.collect::<Vec<_>>()
});
expect_test::expect![[r#"
[
[
Some(
"image-b",
),
Some(
"image-a",
),
],
[
Some(
"image-a",
),
Some(
"image-b",
),
],
[],
]
"#]]
.assert_debug_eq(&uuids);
}
fn image_url_part() -> MediaContentPart {
MediaContentPart::ImageUrl {
url: "https://example.com/image.png".to_string(),
rust/src/chat/src/renderer/deepseek.rs
+81 -31
View File
@@ -16,6 +16,7 @@ use serde_json_fmt::JsonFormat;
use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER;
use super::MediaPartSource;
use crate::error::{Error, Result};
use crate::request::{
ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort,
@@ -87,6 +88,14 @@ struct RenderedToolSchema<'a> {
/// Render one chat request into the final prompt string.
pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result<String> {
render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt)
}
/// Render one chat request and record media in placeholder order.
pub(super) fn render_request_with_media_order(
request: &ChatRequest,
dialect: DsDialect,
) -> Result<(String, Vec<MediaPartSource>)> {
let (thinking_mode, reasoning_effort_prompt) = match dialect {
DsDialect::V4 => {
resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))?
@@ -112,6 +121,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
dialect,
);
let mut out = String::from(BOS_TOKEN);
let mut media_order = Vec::new();
if dialect == DsDialect::V41
&& (thinking_mode == ThinkingMode::Thinking
|| synthetic_tool_system
@@ -126,7 +136,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
let mut request_tools_attached = false;
let mut render_index = 0isize;
if synthetic_tool_system {
render_system_message(&mut out, None, request_tools, dialect)?;
render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?;
request_tools_attached = true;
render_index += 1;
}
@@ -157,14 +167,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
} else {
&[]
};
render_system_message(&mut out, Some(content), tools, dialect)?;
render_system_message(
&mut out,
Some((message_index, content)),
tools,
dialect,
&mut media_order,
)?;
}
ChatMessage::Developer { content, tools } => {
render_developer_message(
&mut out,
message_index,
content,
tools.as_deref().unwrap_or(&[]),
dialect,
&mut media_order,
)?;
}
ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => {
@@ -174,6 +192,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
message_index,
&last_tool_call_order,
dialect,
&mut media_order,
)?;
}
ChatMessage::Assistant { content } => {
@@ -205,25 +224,34 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul
}
}
if (is_user_like_entry(message, current_render_index as usize, dialect)
|| (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. })))
&& next_rendered_entry_is_assistant_or_end(
if is_user_like_entry(message, current_render_index as usize, dialect)
|| (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. }))
{
// Skip the opener only at the last rendered entry when
// add_generation_prompt is off. Inspect the next rendered
// entry, not the raw index: V4 merges trailing tool results.
let should_write_transition = match next_rendered_entry(
request.messages.as_slice(),
message_index,
drop_historical_developers,
last_user_like_message_index,
)
{
write_assistant_transition(
&mut out,
thinking_mode,
drop_thinking,
current_render_index >= last_user_render_index,
);
) {
Some(ChatMessage::Assistant { .. }) => true,
None => request.chat_options.add_generation_prompt(),
Some(_) => false,
};
if should_write_transition {
write_assistant_transition(
&mut out,
thinking_mode,
drop_thinking,
current_render_index >= last_user_render_index,
);
}
}
}
Ok(out)
Ok((out, media_order))
}
/// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer
@@ -401,14 +429,17 @@ fn is_historical_developer(
&& last_user_like_message_index.is_some_and(|last_index| message_index < last_index)
}
/// Return whether the next rendered entry is assistant, or there is no next
/// entry.
fn next_rendered_entry_is_assistant_or_end(
/// Return the next rendered message after `message_index`.
///
/// Skips merged user-content tails and dropped historical developers so the
/// caller can tell a following assistant turn apart from the end of the
/// conversation.
fn next_rendered_entry(
messages: &[ChatMessage],
message_index: usize,
drop_historical_developers: bool,
last_user_like_message_index: Option<usize>,
) -> bool {
) -> Option<&ChatMessage> {
let mut next_index = message_index + 1;
while next_index < messages.len()
&& (is_following_user_content(messages, next_index)
@@ -422,10 +453,7 @@ fn next_rendered_entry_is_assistant_or_end(
next_index += 1;
}
messages
.get(next_index)
.map(|message| matches!(message, ChatMessage::Assistant { .. }))
.unwrap_or(true)
messages.get(next_index)
}
/// Render the tool preamble shown to the model for one DeepSeek dialect.
@@ -488,12 +516,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
/// Render a system turn, optionally followed by the V4 tool preamble.
fn render_system_message(
out: &mut String,
content: Option<&ChatContent>,
content: Option<(usize, &ChatContent)>,
tools: &[ChatTool],
dialect: DsDialect,
media_order: &mut Vec<MediaPartSource>,
) -> Result<()> {
if let Some(content) = content {
write_chat_content(out, content, dialect)?;
if let Some((message_index, content)) = content {
write_chat_content(out, content, dialect, message_index, media_order)?;
}
if !tools.is_empty() {
out.push_str("\n\n");
@@ -505,9 +534,11 @@ fn render_system_message(
/// Developer messages are rendered as user-like turns with optional tools.
fn render_developer_message(
out: &mut String,
message_index: usize,
content: &ChatContent,
tools: &[ChatTool],
dialect: DsDialect,
media_order: &mut Vec<MediaPartSource>,
) -> Result<()> {
if content.is_empty() {
return Err(Error::ChatTemplate(
@@ -516,7 +547,7 @@ fn render_developer_message(
}
out.push_str(USER_SP_TOKEN);
write_chat_content(out, content, dialect)?;
write_chat_content(out, content, dialect, message_index, media_order)?;
if !tools.is_empty() {
out.push_str("\n\n");
render_tools(out, tools, dialect)?;
@@ -531,6 +562,7 @@ fn render_user_content_block(
message_index: usize,
tool_call_order: &HashMap<String, usize>,
dialect: DsDialect,
media_order: &mut Vec<MediaPartSource>,
) -> Result<()> {
let (block_start, block_end) = user_content_block_bounds(messages, message_index);
let mut sorted_tool_indices =
@@ -542,7 +574,9 @@ fn render_user_content_block(
out.push_str("\n\n");
}
match &messages[message_index] {
ChatMessage::User { content } => write_chat_content(out, content, dialect)?,
ChatMessage::User { content } => {
write_chat_content(out, content, dialect, message_index, media_order)?;
}
ChatMessage::ToolResponse { .. } => {
let sorted_index = sorted_tool_indices
.next()
@@ -550,7 +584,7 @@ fn render_user_content_block(
let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else {
unreachable!("sorted tool response index should reference a tool message");
};
write_tool_result(out, content, dialect)?;
write_tool_result(out, content, dialect, sorted_index, media_order)?;
}
_ => unreachable!("user content block should only contain user content messages"),
}
@@ -597,9 +631,15 @@ fn sorted_tool_response_indices(
}
/// Render one tool response payload inside a V4 `<tool_result>` block.
fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
fn write_tool_result(
out: &mut String,
content: &ChatContent,
dialect: DsDialect,
message_index: usize,
media_order: &mut Vec<MediaPartSource>,
) -> Result<()> {
out.push_str("<tool_result>");
write_chat_content(out, content, dialect)?;
write_chat_content(out, content, dialect, message_index, media_order)?;
out.push_str("</tool_result>");
Ok(())
}
@@ -727,7 +767,13 @@ fn encode_arguments_to_dsml(
/// The V4.1 dialect inlines the image placeholder at each image part's
/// position unconditionally (matching the Python encoding's
/// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts.
fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> {
fn write_chat_content(
out: &mut String,
content: &ChatContent,
dialect: DsDialect,
message_index: usize,
media_order: &mut Vec<MediaPartSource>,
) -> Result<()> {
match content {
ChatContent::Text(text) => out.push_str(text),
ChatContent::Parts(parts) => {
@@ -738,6 +784,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec
match part {
ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => {
out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER);
media_order.push(MediaPartSource {
message_index,
content_part_index: index,
});
}
_ => out.push_str(part.as_text()?),
}
rust/src/chat/src/renderer/deepseek_v32/encoding.rs
+14 -4
View File
@@ -61,6 +61,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
find_last_user_render_index(request.messages.as_slice(), render_offset);
let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice());
let continue_final_message = request.chat_options.continue_final_message();
let add_generation_prompt = request.chat_options.add_generation_prompt();
let mut prompt = String::from(BOS_TOKEN);
if !initial_tools.is_empty() {
@@ -79,6 +80,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
thinking_mode,
drop_thinking,
continue_final_message,
add_generation_prompt,
)?;
}
@@ -110,11 +112,13 @@ fn render_message(
thinking_mode: ThinkingMode,
drop_thinking: bool,
continue_final_message: bool,
add_generation_prompt: bool,
) -> Result<()> {
let render_index = message_index as isize + render_offset;
let opens_thinking = render_index == last_user_render_index;
let after_last_user_turn = render_index > last_user_render_index;
let after_or_at_last_user_turn = render_index >= last_user_render_index;
let add_generation_prompt = message_index + 1 != messages.len() || add_generation_prompt;
match message {
ChatMessage::System { content } => render_system_message(out, Some(content), &[]),
@@ -123,11 +127,13 @@ fn render_message(
content,
tools.as_deref().unwrap_or(&[]),
thinking_mode == ThinkingMode::Thinking && opens_thinking,
add_generation_prompt,
),
ChatMessage::User { content } => render_user_message(
out,
content,
thinking_mode == ThinkingMode::Thinking && opens_thinking,
add_generation_prompt,
),
ChatMessage::Assistant { content } => render_assistant_message(
out,
@@ -195,6 +201,7 @@ fn render_developer_message(
content: &ChatContent,
tools: &[ChatTool],
opens_thinking: bool,
add_generation_prompt: bool,
) -> Result<()> {
if content.is_empty() {
return Err(Error::ChatTemplate(
@@ -209,7 +216,7 @@ fn render_developer_message(
}
out.push_str("\n\n# The user's message is: ");
write_chat_content(out, content)?;
write_user_like_suffix(out, opens_thinking);
write_user_like_suffix(out, opens_thinking, add_generation_prompt);
Ok(())
}
@@ -219,17 +226,20 @@ fn render_user_message(
out: &mut String,
content: &ChatContent,
opens_thinking: bool,
add_generation_prompt: bool,
) -> Result<()> {
out.push_str("<｜User｜>");
write_chat_content(out, content)?;
write_user_like_suffix(out, opens_thinking);
write_user_like_suffix(out, opens_thinking, add_generation_prompt);
Ok(())
}
/// Shared trailing wrapper used by both real user turns and native developer
/// turns after their content has already been written.
// TODO: respect `add_generation_prompt` option
fn write_user_like_suffix(out: &mut String, opens_thinking: bool) {
fn write_user_like_suffix(out: &mut String, opens_thinking: bool, add_generation_prompt: bool) {
if !add_generation_prompt {
return;
}
out.push_str("<｜Assistant｜>");
if opens_thinking {
out.push_str(THINKING_START_TOKEN);
rust/src/chat/src/renderer/deepseek_v32/mod.rs
+1
View File
@@ -26,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer {
Ok(RenderedPrompt {
prompt: Prompt::Text(encoding::render_request(request)?),
media_order: None,
effective_template_kwargs: request_template_kwargs(request),
})
}
rust/src/chat/src/renderer/deepseek_v32/tests.rs
+38
View File
@@ -281,6 +281,44 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() {
.assert_eq(&error.to_report_string());
}
#[test]
fn last_user_turn_appends_generation_prompt_by_default() {
let request = thinking_request(vec![ChatMessage::user("Hello")]);
let rendered = render_request(&request);
assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜><think>"));
}
#[test]
fn last_user_turn_omits_generation_prompt_when_disabled() {
let mut request = thinking_request(vec![ChatMessage::user("Hello")]);
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_request(&request);
assert!(rendered.ends_with("<｜User｜>Hello"));
assert!(!rendered.contains("<｜Assistant｜>"));
assert!(!rendered.contains("<think>"));
}
#[test]
fn user_then_assistant_still_writes_suffix_without_generation_prompt() {
let mut request = ChatRequest {
messages: vec![
ChatMessage::user("Hello"),
ChatMessage::assistant_text("Hi there."),
],
..ChatRequest::for_test()
};
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_request(&request);
expect!["<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>Hi there.<｜end▁of▁sentence｜>"]
.assert_eq(&rendered);
}
#[test]
fn continue_final_assistant_omits_final_eos() {
let mut request = ChatRequest {
rust/src/chat/src/renderer/deepseek_v4/mod.rs
+1
View File
@@ -24,6 +24,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer {
Ok(RenderedPrompt {
prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?),
media_order: None,
effective_template_kwargs: request_template_kwargs(request),
})
}
rust/src/chat/src/renderer/deepseek_v4/tests.rs
+91
View File
@@ -633,3 +633,94 @@ fn continue_final_assistant_omits_final_eos() {
expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜></think>partial answer"]
.assert_eq(&rendered);
}
#[test]
fn last_user_turn_appends_generation_prompt_by_default() {
let request = thinking_request(vec![ChatMessage::user("Hello")]);
let rendered = render_request(&request);
assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜><think>"));
}
#[test]
fn last_user_turn_omits_generation_prompt_when_disabled() {
let mut request = thinking_request(vec![ChatMessage::user("Hello")]);
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_request(&request);
assert!(rendered.ends_with("<｜User｜>Hello"));
assert!(!rendered.contains("<｜Assistant｜>"));
assert!(!rendered.contains("<think>"));
}
#[test]
fn user_then_assistant_still_writes_suffix_without_generation_prompt() {
let mut request = ChatRequest {
messages: vec![
ChatMessage::user("Hello"),
ChatMessage::assistant_text("Hi there."),
],
..ChatRequest::for_test()
};
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
request
.chat_options
.template_kwargs
.insert("thinking".to_string(), Value::Bool(false));
let rendered = render_request(&request);
expect![
"<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>Hi there.<｜end▁of▁sentence｜>"
]
.assert_eq(&rendered);
}
#[test]
fn trailing_system_omits_generation_prompt_when_disabled() {
let mut request = ChatRequest {
messages: vec![
ChatMessage::system("rules"),
ChatMessage::user("question"),
ChatMessage::system("final rules"),
],
..ChatRequest::for_test()
};
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
request
.chat_options
.template_kwargs
.insert("enable_thinking".to_string(), Value::Bool(false));
request.chat_options.reasoning_effort = Some(ReasoningEffort::Low);
let rendered = render_request(&request);
expect!["<｜begin▁of▁sentence｜>rules<｜User｜>questionfinal rules"].assert_eq(&rendered);
}
#[test]
fn last_merged_tool_result_omits_generation_prompt_when_disabled() {
let mut request = thinking_request(vec![
ChatMessage::user("question"),
ChatMessage::tool_response("result", "call-1"),
]);
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_request(&request);
expect!["<｜begin▁of▁sentence｜><｜User｜>question\n\n<tool_result>result</tool_result>"]
.assert_eq(&rendered);
assert!(!rendered.contains("<｜Assistant｜>"));
}
#[test]
fn last_developer_omits_generation_prompt_when_disabled() {
let mut request = thinking_request(vec![ChatMessage::developer("latest instruction", None)]);
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_request(&request);
expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered);
}
rust/src/chat/src/renderer/deepseek_v41/mod.rs
+4 -1
View File
@@ -21,9 +21,12 @@ impl DeepSeekV41ChatRenderer {
impl ChatRenderer for DeepSeekV41ChatRenderer {
fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
request.validate()?;
let (prompt, media_order) =
deepseek::render_request_with_media_order(request, DsDialect::V41)?;
Ok(RenderedPrompt {
prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?),
prompt: Prompt::Text(prompt),
media_order: Some(media_order),
effective_template_kwargs: request_template_kwargs(request),
})
}
rust/src/chat/src/renderer/deepseek_v41/tests.rs
+123 -2
View File
@@ -8,9 +8,11 @@ use serde_json::{Value, json};
use super::DeepSeekV41ChatRenderer;
use crate::ChatRenderer;
use crate::event::AssistantContentBlock;
use crate::event::{AssistantContentBlock, AssistantToolCall};
use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort};
use crate::request::{
ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort,
};
fn render(request: &ChatRequest) -> String {
DeepSeekV41ChatRenderer::new()
@@ -83,6 +85,82 @@ fn inlines_image_placeholder_at_image_part_positions() {
);
}
#[test]
fn media_order_follows_reordered_tool_results() {
let renderer = DeepSeekV41ChatRenderer::new();
let request = ChatRequest {
messages: vec![
ChatMessage::assistant_blocks(vec![
AssistantContentBlock::ToolCall(AssistantToolCall {
id: "call-a".to_string(),
name: "tool-a".to_string(),
arguments: "{}".to_string(),
}),
AssistantContentBlock::ToolCall(AssistantToolCall {
id: "call-b".to_string(),
name: "tool-b".to_string(),
arguments: "{}".to_string(),
}),
]),
ChatMessage::tool_response(
vec![ChatContentPart::text("result-b"), image_part("image-b")],
"call-b",
),
ChatMessage::tool_response(
vec![ChatContentPart::text("result-a"), image_part("image-a")],
"call-a",
),
],
..ChatRequest::for_test()
};
let rendered = renderer.render(&request).unwrap();
let media_order = rendered.media_order.unwrap();
let prompt = rendered.prompt.into_text().unwrap();
assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
expect![[r#"
[
MediaPartSource {
message_index: 2,
content_part_index: 1,
},
MediaPartSource {
message_index: 1,
content_part_index: 1,
},
]
"#]]
.assert_debug_eq(&media_order);
}
#[test]
fn media_order_excludes_dropped_historical_developer_content() {
let renderer = DeepSeekV41ChatRenderer::new();
let request = ChatRequest {
messages: vec![
ChatMessage::developer(vec![image_part("historical-image")], None),
ChatMessage::user("question"),
],
..ChatRequest::for_test()
};
let rendered = renderer.render(&request).unwrap();
let media_order = rendered.media_order.unwrap();
let prompt = rendered.prompt.into_text().unwrap();
assert!(!prompt.contains("<｜deepseek_image｜>"));
assert!(media_order.is_empty());
}
fn image_part(uuid: &str) -> ChatContentPart {
ChatContentPart::ImageUrl {
image_url: format!("data:image/png;base64,{uuid}"),
detail: None,
uuid: Some(uuid.to_string()),
}
}
#[test]
fn maps_reasoning_effort_to_reference_numeric_budget() {
for (effort, budget) in [
@@ -205,6 +283,49 @@ fn mid_system_advances_last_user_boundary_and_drops_old_reasoning() {
.assert_eq(&render(&request));
}
#[test]
fn last_user_turn_omits_generation_prompt_when_disabled() {
let mut request = request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question"]
.assert_eq(&render(&request));
}
#[test]
fn last_mid_system_omits_generation_prompt_when_disabled() {
let mut request = ChatRequest {
messages: vec![
ChatMessage::system("sys"),
ChatMessage::user("q1"),
ChatMessage::assistant_blocks(vec![
AssistantContentBlock::Reasoning { text: "r1".into() },
AssistantContentBlock::Text { text: "a1".into() },
]),
ChatMessage::developer("old policy", None),
ChatMessage::system("mid sys"),
],
..ChatRequest::for_test()
};
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
request
.chat_options
.template_kwargs
.insert("reasoning_effort".into(), json!(88));
expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜></think>a1<｜end▁of▁sentence｜><｜System｜>mid sys"]
.assert_eq(&render(&request));
}
#[test]
fn chat_mode_omits_generation_prompt_when_disabled() {
let mut request = request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
request.chat_options.template_kwargs.insert("thinking".into(), json!(false));
expect!["<｜begin▁of▁sentence｜><｜User｜>question"].assert_eq(&render(&request));
}
#[test]
fn joins_text_parts_with_reference_separator_and_preserves_literal_tags() {
let mut request = ChatRequest {
rust/src/chat/src/renderer/harmony/mod.rs
+1
View File
@@ -116,6 +116,7 @@ impl ChatRenderer for HarmonyChatRenderer {
fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
Ok(RenderedPrompt {
prompt: Prompt::TokenIds(self.render_token_ids(request)?),
media_order: None,
effective_template_kwargs: request_template_kwargs(request),
})
}
rust/src/chat/src/renderer/hf/mod.rs
+1
View File
@@ -217,6 +217,7 @@ impl HfChatRenderer {
Ok(RenderedPrompt {
prompt: Prompt::Text(prompt),
media_order: None,
effective_template_kwargs,
})
}
rust/src/chat/src/renderer/inkling/mod.rs
+1
View File
@@ -346,6 +346,7 @@ impl ChatRenderer for InklingChatRenderer {
Ok(RenderedPrompt {
prompt: Prompt::TokenIds(out),
media_order: None,
effective_template_kwargs,
})
}
rust/src/chat/src/renderer/kimi_k3/encoding.rs
+37 -14
View File
@@ -10,6 +10,7 @@ use std::collections::HashMap;
use serde_json::{Map, Value, json};
use vllm_tokenizer::Tokenizer;
use super::super::MediaPartSource;
use crate::error::{Error, Result};
use crate::request::{
ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice,
@@ -29,6 +30,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"];
pub(super) struct K3TokenWriter<'a> {
tokenizer: &'a dyn Tokenizer,
token_ids: Vec<u32>,
media_order: Vec<MediaPartSource>,
}
impl<'a> K3TokenWriter<'a> {
@@ -36,6 +38,7 @@ impl<'a> K3TokenWriter<'a> {
Self {
tokenizer,
token_ids: Vec::new(),
media_order: Vec::new(),
}
}
@@ -55,13 +58,25 @@ impl<'a> K3TokenWriter<'a> {
Ok(())
}
pub(super) fn finish(self) -> Vec<u32> {
self.token_ids
fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> {
self.control(IMAGE_PLACEHOLDER)?;
self.media_order.push(MediaPartSource {
message_index,
content_part_index,
});
Ok(())
}
pub(super) fn finish(self) -> (Vec<u32>, Vec<MediaPartSource>) {
(self.token_ids, self.media_order)
}
}
/// Render and tokenize one chat request using K3's segment-aware contract.
pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result<Vec<u32>> {
/// Render and tokenize one chat request while recording media in placeholder order.
pub(super) fn render_request_with_media_order(
request: &ChatRequest,
tokenizer: &dyn Tokenizer,
) -> Result<(Vec<u32>, Vec<MediaPartSource>)> {
let thinking = thinking_enabled(request)?;
let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?;
let tools = request_tools(request);
@@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
resolved.sort_by_key(|item| (item.0, item.1));
}
for (position, _, content, name) in resolved {
for (position, message_index, content, name) in resolved {
let tool_name = name.as_deref().ok_or_else(|| {
Error::ChatTemplate(
"Kimi K3 tool messages need a resolvable tool name: \
@@ -135,7 +150,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
.to_string(),
)
})?;
write_tool_message(out, tool_name, position, &content)?;
write_tool_message(out, tool_name, position, message_index, &content)?;
}
Ok(())
};
@@ -161,14 +176,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -
} if !local_tools.is_empty() => {
write_tool_declare(&mut out, local_tools, true)?;
if !content_is_empty(content) {
write_role_message(&mut out, "system", None, content)?;
write_role_message(&mut out, "system", None, message_index, content)?;
}
}
ChatMessage::System { content } | ChatMessage::Developer { content, .. } => {
write_role_message(&mut out, "system", None, content)?;
write_role_message(&mut out, "system", None, message_index, content)?;
}
ChatMessage::User { content } => {
write_role_message(&mut out, "user", None, content)?;
write_role_message(&mut out, "user", None, message_index, content)?;
}
ChatMessage::Assistant { content } => {
tool_call_id_index.clear();
@@ -364,6 +379,7 @@ fn write_role_message(
out: &mut K3TokenWriter<'_>,
role: &str,
name: Option<&str>,
message_index: usize,
content: &ChatContent,
) -> Result<()> {
let mut attrs = vec![("role", role.to_string())];
@@ -372,7 +388,7 @@ fn write_role_message(
}
let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect();
write_open_tag(out, "message", &attr_refs)?;
write_content(out, content)?;
write_content(out, message_index, content)?;
write_close_tag(out, "message")?;
out.control(END_OF_MSG)
}
@@ -381,6 +397,7 @@ fn write_tool_message(
out: &mut K3TokenWriter<'_>,
tool_name: &str,
index: usize,
message_index: usize,
content: &ChatContent,
) -> Result<()> {
let index_str = index.to_string();
@@ -389,7 +406,7 @@ fn write_tool_message(
"message",
&[("role", "tool"), ("tool", tool_name), ("index", &index_str)],
)?;
write_content(out, content)?;
write_content(out, message_index, content)?;
write_close_tag(out, "message")?;
out.control(END_OF_MSG)
}
@@ -471,14 +488,20 @@ fn write_assistant_tool_call(
write_close_tag(out, "call")
}
fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> {
fn write_content(
out: &mut K3TokenWriter<'_>,
message_index: usize,
content: &ChatContent,
) -> Result<()> {
match content {
ChatContent::Text(text) => write_text_with_images(out, text),
ChatContent::Parts(parts) => {
for part in parts {
for (content_part_index, part) in parts.iter().enumerate() {
match part {
ChatContentPart::Text { text } => write_text_with_images(out, text)?,
ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?,
ChatContentPart::ImageUrl { .. } => {
out.image(message_index, content_part_index)?;
}
ChatContentPart::VideoUrl { .. } => {
return Err(Error::UnsupportedMultimodalContent("video_url"));
}
rust/src/chat/src/renderer/kimi_k3/mod.rs
+4 -1
View File
@@ -30,9 +30,12 @@ impl KimiK3ChatRenderer {
impl ChatRenderer for KimiK3ChatRenderer {
fn render(&self, request: &ChatRequest) -> Result<RenderedPrompt> {
request.validate()?;
let (token_ids, media_order) =
encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?;
Ok(RenderedPrompt {
prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?),
prompt: Prompt::TokenIds(token_ids),
media_order: Some(media_order),
effective_template_kwargs: request_template_kwargs(request),
})
}
rust/src/chat/src/renderer/kimi_k3/tests.rs
+59
View File
@@ -284,6 +284,65 @@ fn partial_tool_results_keep_assistant_call_position() {
);
}
#[test]
fn media_order_follows_reordered_tool_results() {
let mut request = crate::request::ChatRequest::for_test();
request.messages = vec![
ChatMessage::assistant_blocks(vec![
AssistantContentBlock::ToolCall(AssistantToolCall {
id: "call-a".to_string(),
name: "tool-a".to_string(),
arguments: "{}".to_string(),
}),
AssistantContentBlock::ToolCall(AssistantToolCall {
id: "call-b".to_string(),
name: "tool-b".to_string(),
arguments: "{}".to_string(),
}),
]),
ChatMessage::tool_response(
vec![ChatContentPart::text("result-b"), image_part("image-b")],
"call-b",
),
ChatMessage::tool_response(
vec![ChatContentPart::text("result-a"), image_part("image-a")],
"call-a",
),
];
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer()));
let rendered = renderer.render(&request).unwrap();
let media_order = rendered.media_order.unwrap();
let Prompt::TokenIds(token_ids) = rendered.prompt else {
panic!("kimi k3 renderer should return token IDs")
};
let prompt = test_tokenizer().decode(&token_ids, false).unwrap();
assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap());
expect![[r#"
[
MediaPartSource {
message_index: 2,
content_part_index: 1,
},
MediaPartSource {
message_index: 1,
content_part_index: 1,
},
]
"#]]
.assert_debug_eq(&media_order);
}
fn image_part(uuid: &str) -> ChatContentPart {
ChatContentPart::ImageUrl {
image_url: format!("data:image/png;base64,{uuid}"),
detail: None,
uuid: Some(uuid.to_string()),
}
}
#[test]
fn defaults_thinking_effort_to_max() {
let rendered = render_request(&crate::request::ChatRequest::for_test());
rust/src/chat/src/renderer/mod.rs
+13
View File
@@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer;
pub use kimi_k3::KimiK3ChatRenderer;
pub use selection::RendererSelection;
/// Location of one multimodal content part in the source chat request.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MediaPartSource {
/// Zero-based message index in the source request.
pub message_index: usize,
/// Zero-based content-part index within the source message.
pub content_part_index: usize,
}
/// Rendered chat prompt submitted to the text backend.
#[derive(Debug, Clone, PartialEq)]
pub struct RenderedPrompt {
/// The rendered prompt, either as text or already tokenized.
pub prompt: Prompt,
/// Media sources in rendered placeholder order.
/// `None` uses the message and content-part order in the request (e.g. Jinja).
/// `Some` uses exactly the listed sources; an empty list omits all media.
pub media_order: Option<Vec<MediaPartSource>>,
/// Effective chat-template kwargs visible to the renderer after applying
/// server defaults, request overrides, and typed reasoning controls.
pub effective_template_kwargs: HashMap<String, Value>,
rust/src/chat/tests/chat.rs
+1
View File
@@ -250,6 +250,7 @@ impl ChatRenderer for FakeChatBackend {
Ok(RenderedPrompt {
prompt: Prompt::Text(prompt),
media_order: None,
effective_template_kwargs: request.chat_options.template_kwargs.clone(),
})
}
rust/src/cmd/Cargo.toml
+1
View File
@@ -33,6 +33,7 @@ vllm-chat.workspace = true
vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true
vllm-server.workspace = true
vllm-text.workspace = true
vllm-tracing.workspace = true
[dev-dependencies]
rust/src/cmd/src/cli.rs
+27
View File
@@ -33,6 +33,7 @@ use vllm_server::{
DEFAULT_KEEP_ALIVE_TIMEOUT, HttpListenerMode, LoraModulePath, ParserSelection, RenderConfig,
RendererSelection,
};
use vllm_text::backend::hf::HfOverrides;
use crate::cli::ssl::SslArgs;
use crate::cli::unsupported::UnsupportedArgs;
@@ -113,6 +114,9 @@ pub struct RenderArgs {
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
#[arg(long)]
revision: Option<String>,
/// JSON Merge Patch (RFC 7396) for config.json; null removes a field.
#[arg(long, value_parser = parse_json::<HfOverrides>, default_value = "{}", value_name = "JSON")]
hf_overrides: HfOverrides,
/// HTTP bind host.
#[arg(long, default_value = "127.0.0.1")]
host: String,
@@ -161,6 +165,7 @@ impl RenderArgs {
RenderConfig {
model: self.model,
revision: self.revision,
hf_overrides: self.hf_overrides,
served_model_name: self.served_model_name,
host: self.host,
port: self.port,
@@ -204,6 +209,12 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub revision: Option<String>,
/// JSON Merge Patch (RFC 7396) for config.json; null removes a field.
/// Objects merge recursively and arrays/scalars replace existing values.
#[arg(long, value_parser = parse_json::<HfOverrides>, default_value = "{}", value_name = "JSON")]
#[serde(default)]
pub hf_overrides: HfOverrides,
/// The source of generation-config sampling defaults. `"auto"` loads the
/// model's defaults, while `"vllm"` uses vLLM's neutral defaults.
#[arg(long, default_value_t)]
@@ -329,6 +340,15 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub enable_request_id_headers: bool,
/// Register the scale-out `/inference/v1/generate` endpoint.
#[arg(
long,
default_missing_value = "true",
num_args = 0..=1
)]
#[serde(default)]
pub enable_scale_out: bool,
/// If provided, the server will require one of these keys to be presented
/// in the Authorization header.
#[educe(Debug(ignore))]
@@ -491,6 +511,7 @@ impl SharedRuntimeArgs {
},
model: self.model,
revision: self.revision,
hf_overrides: self.hf_overrides,
generation_config: self.generation_config,
served_model_name: self.served_model_name,
listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd },
@@ -547,6 +568,7 @@ impl SharedRuntimeArgs {
coordinator_mode: CoordinatorMode::MaybeInProc,
model: self.model,
revision: self.revision,
hf_overrides: self.hf_overrides,
generation_config: self.generation_config,
served_model_name: self.served_model_name,
listener_mode,
@@ -577,6 +599,7 @@ impl SharedRuntimeArgs {
enable_log_requests: self.enable_log_requests,
enable_prompt_tokens_details: self.enable_prompt_tokens_details,
enable_request_id_headers: self.enable_request_id_headers,
enable_scale_out: self.enable_scale_out,
}
}
@@ -751,6 +774,9 @@ impl ServeArgs {
let reasoning_parser =
effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model);
let profiler_config = self.runtime.profiler_config_json();
let hf_overrides = (!self.runtime.hf_overrides.is_empty()).then(|| {
serde_json::to_string(&self.runtime.hf_overrides).expect("JSON object serializes")
});
self.managed_engine.clone().into_config(
self.runtime.model.clone(),
@@ -763,6 +789,7 @@ impl ServeArgs {
self.runtime.shutdown_timeout,
handshake_port,
self.runtime.limit_mm_per_prompt_json(),
hf_overrides,
)
}
}
rust/src/cmd/src/cli/tests.rs
+91
View File
@@ -177,6 +177,9 @@ fn serve_args_forward_python_flags_with_separator() {
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: Auto,
@@ -195,6 +198,7 @@ fn serve_args_forward_python_flags_with_separator() {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
disable_log_stats: false,
served_model_name: [],
allowed_origins: JsonStringList(
@@ -1012,6 +1016,9 @@ fn frontend_args_accept_json() {
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: None,
@@ -1030,6 +1037,7 @@ fn frontend_args_accept_json() {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
disable_log_stats: false,
served_model_name: [],
allowed_origins: JsonStringList(
@@ -1118,6 +1126,73 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() {
assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B");
}
#[test]
fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() {
let patch = r#"{"text_config":{"rope_parameters":{"factor":4}},"sliding_window":null}"#;
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--hf-overrides",
patch,
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args")
};
let expected: serde_json::Value = serde_json::from_str(patch).unwrap();
assert_eq!(
serde_json::to_value(&args.runtime.hf_overrides).unwrap(),
expected
);
let engine = args.to_managed_engine_config(1234);
let index = engine.python_args.iter().position(|arg| arg == "--hf-overrides").unwrap();
assert_eq!(
serde_json::from_str::<serde_json::Value>(&engine.python_args[index + 1]).unwrap(),
expected
);
assert!(!engine.python_args.iter().any(|arg| arg == "--hf-config-path"));
let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string());
assert_eq!(serde_json::to_value(config.hf_overrides).unwrap(), expected);
let payload = serde_json::json!({"model_tag":"Qwen/Qwen3-0.6B", "hf_overrides":expected});
let cli = Cli::try_parse_from([
"vllm-rs",
"frontend",
"--listen-fd",
"3",
"--input-address",
"ipc:///tmp/input.sock",
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
&payload.to_string(),
])
.unwrap();
let Command::Frontend(args) = cli.command else {
panic!("expected frontend args")
};
assert_eq!(
serde_json::to_value(args.into_config().hf_overrides).unwrap(),
expected
);
}
#[test]
fn hf_overrides_reject_non_object_inputs() {
for patch in ["null", "[]", "1", r#""callable""#] {
let error = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--hf-overrides",
patch,
])
.unwrap_err();
assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation);
}
}
#[test]
fn frontend_args_json_accepts_supported_non_default_fields() {
let cli = Cli::try_parse_from([
@@ -1611,6 +1686,9 @@ fn serve_args_accept_handshake_aliases() {
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: Auto,
@@ -1629,6 +1707,7 @@ fn serve_args_accept_handshake_aliases() {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
disable_log_stats: false,
served_model_name: [],
allowed_origins: JsonStringList(
@@ -1760,6 +1839,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
served_model_name: [],
listener_mode: BindTcp {
@@ -1780,6 +1862,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
},
cors: CorsConfig {
allow_origins: [
@@ -1848,6 +1931,9 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
served_model_name: [],
listener_mode: BindTcp {
@@ -1868,6 +1954,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
},
cors: CorsConfig {
allow_origins: [
@@ -1958,6 +2045,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
},
model: "Qwen/Qwen3-0.6B",
revision: None,
hf_overrides: HfOverrides(
{},
),
generation_config: Auto,
served_model_name: [],
listener_mode: InheritedFd {
@@ -1977,6 +2067,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
enable_log_requests: false,
enable_prompt_tokens_details: false,
enable_request_id_headers: false,
enable_scale_out: false,
},
cors: CorsConfig {
allow_origins: [
rust/src/cmd/src/cli/unsupported.rs
-5
View File
@@ -235,11 +235,6 @@ pub struct EngineUnsupportedArgs {
#[arg(long, default_missing_value = "true", num_args = 0..=1)]
pub hf_token: Option<Unsupported>,
/// If a dictionary, contains arguments to be forwarded to the Hugging Face
/// config. If a callable, it is called to update the HuggingFace config.
#[arg(long)]
pub hf_overrides: Option<Unsupported>,
/// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If
/// used with `--generation-config auto`, the override parameters will be
/// merged with the default config from the model. If used with
rust/src/engine-core-client/src/client/imp.rs
+25 -1
View File
@@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
use crate::client::stream::EngineCoreStreamOutput;
use crate::client::{AbortCause, AbortRequest};
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder};
use crate::protocol::encode_msgpack;
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
@@ -47,6 +47,7 @@ pub(crate) struct ClientInner {
/// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
msgpack_zero_copy_threshold: usize,
scheduler_stats_recorder: SchedulerStatsRecorder,
iteration_metrics: BTreeMap<u32, IterationMetricHandles>,
request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>,
@@ -64,12 +65,21 @@ impl ClientInner {
) -> Self {
let scheduler_stats_recorder =
SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
let iteration_metrics = engines
.iter()
.filter_map(|engine| {
let engine = engine.engine_id.engine_index()?;
let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine);
Some((engine, handles))
})
.collect();
Self {
input_send,
handle,
model_name,
msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
scheduler_stats_recorder,
iteration_metrics,
request_reg: Mutex::new(RequestRegistry::new(engines)),
utility_reg: Mutex::new(UtilityRegistry::default()),
health_error: ArcSwapOption::empty(),
@@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop(
match outputs {
EngineCoreOutputs::RequestBatch(batch) => {
let has_outputs = !batch.outputs.is_empty();
let mut iteration_tokens = 0_u64;
let senders = inner.take_senders_for_outputs(&batch.outputs);
for (output, sender) in batch.outputs.into_iter().zip(senders) {
let request_id = output.request_id.clone();
@@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop(
continue;
};
iteration_tokens += output.new_token_ids.len() as u64;
if let Some(prefill) = &output.prefill_stats {
// The engine emits prefill_stats once, on the first output.
iteration_tokens += u64::from(prefill.num_computed_tokens);
}
let wrapped_output = EngineCoreStreamOutput {
engine_index: batch.engine_index,
timestamp: batch.timestamp,
@@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop(
}
}
if has_outputs
&& let Some(handles) = inner.iteration_metrics.get(&batch.engine_index)
{
handles.iteration_tokens_total.observe(iteration_tokens as f64);
}
// The sender for normally-finished requests should have already been removed
// from the registry when their final output was dispatched
// above. This serves as a safety net to capture any
rust/src/engine-core-client/src/metrics.rs
+19 -2
View File
@@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
use vllm_metrics::{
EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily,
MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
WaitingReasonLabels,
MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics,
U64Counter, U64Gauge, WaitingReasonLabels,
};
use crate::protocol::stats::{
@@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine;
const WAITING_REASON_CAPACITY: &str = "capacity";
const WAITING_REASON_DEFERRED: &str = "deferred";
/// Cached output-batch metric handles for one model and engine index.
pub(crate) struct IterationMetricHandles {
pub iteration_tokens_total: HistogramMetric,
}
impl IterationMetricHandles {
pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self {
let labels = EngineLabels {
model_name: model_name.to_string(),
engine,
};
Self {
iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels),
}
}
}
/// Cached scheduler-stats metric handles for all engines connected to one
/// frontend client.
pub(crate) struct SchedulerStatsRecorder {
rust/src/engine-core-client/src/protocol/logprobs.rs
+3 -4
View File
@@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TokenLogprob {
pub token_id: u32,
/// Preserves the engine's value, including NaN and infinities.
pub logprob: f32,
/// The sampled/selected token uses its actual vocab rank. Remaining entries
/// use 1-based top-k ranks matching the engine's returned candidate
/// order.
/// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's
/// `(logprobs >= selected_logprob).sum(-1)` counts no matching values.
pub rank: u32,
}
@@ -53,10 +56,6 @@ impl PositionLogprobs {
logprobs.len()
);
}
if sampled_rank == 0 {
bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs");
}
let mut entries = Vec::with_capacity(token_ids.len());
for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() {
let rank = if index == 0 {
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+118
View File
@@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() {
"new_logprobs: zero-column logprobs payload with 2 rows"
);
}
fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value {
let rows = ranks.len();
let cols = ids.len() / rows;
Value::Array(vec![
ndarray_value(
"<i8",
&[rows, cols],
Value::Ext(3, ids.iter().flat_map(|v| v.to_le_bytes()).collect()),
),
ndarray_value(
"<f4",
&[rows, cols],
Value::Ext(3, probs.iter().flat_map(|v| v.to_le_bytes()).collect()),
),
ndarray_value(
"<i8",
&[rows],
Value::Ext(3, ranks.iter().flat_map(|v| v.to_le_bytes()).collect()),
),
Value::Nil,
])
}
#[test]
fn zero_sampled_rank_does_not_fail_frame_mates() {
// A batched EngineCoreOutputs frame where one request's row has rank 0
// must still resolve; the other requests' logprobs stay intact.
let frames = vec![Bytes::from(encode_value(&Value::Array(vec![
Value::from(0),
Value::Array(vec![
Value::Array(vec![
Value::from("req-nan"),
Value::Array(vec![Value::from(7), Value::from(8)]),
logprobs_value(&[1, 2], &[f32::NAN, f32::NEG_INFINITY], &[0]),
logprobs_value(&[7, 8], &[f32::NAN, f32::NEG_INFINITY], &[0]),
Value::Nil,
Value::from(EngineCoreFinishReason::Length as u8),
]),
Value::Array(vec![
Value::from("req-ok"),
Value::Array(vec![Value::from(9)]),
logprobs_value(&[3, 4], &[-3.0, -0.5], &[7]),
Value::Nil,
Value::Nil,
Value::from(EngineCoreFinishReason::Length as u8),
]),
]),
Value::Nil,
Value::from(0.0),
Value::Nil,
Value::Nil,
])))];
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let bad = decoded.outputs[0].new_logprobs.as_ref().unwrap();
let prompt = decoded.outputs[0].new_prompt_logprobs_tensors.as_ref().unwrap();
let good = decoded.outputs[1].new_logprobs.as_ref().unwrap();
expect_test::expect![[r#"
(
Logprobs {
positions: [
PositionLogprobs {
entries: [
TokenLogprob {
token_id: 1,
logprob: NaN,
rank: 0,
},
TokenLogprob {
token_id: 2,
logprob: -inf,
rank: 1,
},
],
},
],
},
Logprobs {
positions: [
PositionLogprobs {
entries: [
TokenLogprob {
token_id: 7,
logprob: NaN,
rank: 0,
},
TokenLogprob {
token_id: 8,
logprob: -inf,
rank: 1,
},
],
},
],
},
Logprobs {
positions: [
PositionLogprobs {
entries: [
TokenLogprob {
token_id: 3,
logprob: -3.0,
rank: 7,
},
TokenLogprob {
token_id: 4,
logprob: -0.5,
rank: 1,
},
],
},
],
},
)
"#]]
.assert_debug_eq(&(&**bad, &**prompt, &**good));
}

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