Compare commits

..
Author SHA1 Message Date
Isotr0pyandkhluu 5a532c7e0b Revert "[Misc] Move pyav and soundfile to common requirements" (#40276)
Co-authored-by: Roger Wang <hey@rogerw.io>
(cherry picked from commit 67eb6083e3)
2026-04-22 02:00:47 -07:00
Wentao Yeandkhluu 2f3555bf53 [Deprecation] Deprecate cprofile and cprofile_context (#39100)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
(cherry picked from commit 301024aa9c)
2026-04-22 02:00:12 -07:00
Jhao-Ting Chenandkhluu 4699f1bf8b fix: clamp NaN/Inf in topk_softmax to prevent duplicate expert IDs (#39391)
Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
(cherry picked from commit 28c222157b)
2026-04-22 01:59:52 -07:00
Nicolò Lucchesiandkhluu 681a6371cc [Bugfix][CI] Fix tests/distributed/test_torchrun_example_moe.py (#40349)
Signed-off-by: NickLucche <nlucches@redhat.com>
(cherry picked from commit 304d5ba1a0)
2026-04-22 01:59:37 -07:00
TJianandkhluu e86d349053 [ROCm] [Wheel] [Bugfix] [Critical] Remove any packages installed from github from rocm.txt e.g fastsafetensors as it is incompatible with uv pip (#40461)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
(cherry picked from commit 583e6f2226)
2026-04-22 01:58:51 -07:00
683 changed files with 14319 additions and 51471 deletions
@@ -46,7 +46,7 @@ echo "Image not found, proceeding with build..."
# --- CUDA 13.0 for nightly builds ---
# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
NIGHTLY_CUDA_VERSION="13.0.2"
NIGHTLY_CUDA_VERSION="13.0.0"
NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
-21
View File
@@ -1,21 +0,0 @@
group: Engine Intel
depends_on:
- image-build-xpu
steps:
- label: Engine (1 GPU)
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/v1/engine/
- tests/v1/engine/
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
-21
View File
@@ -1,21 +0,0 @@
group: Kernels Intel
depends_on:
- image-build-xpu
steps:
- label: vLLM IR Tests
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/ir
- vllm/kernels
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s kernels/ir'
-130
View File
@@ -1,130 +0,0 @@
group: LoRA Intel
depends_on:
- image-build-xpu
steps:
- label: LoRA Runtime + Utils
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s lora/test_layers.py &&
pytest -v -s lora/test_lora_checkpoints.py &&
(pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
pytest -v -s lora/test_lora_huggingface.py &&
pytest -v -s lora/test_lora_manager.py &&
pytest -v -s lora/test_lora_utils.py &&
pytest -v -s lora/test_peft_helper.py &&
pytest -v -s lora/test_resolver.py &&
pytest -v -s lora/test_utils.py &&
(pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
(pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
- label: LoRA Fused/MoE Kernels
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s lora/test_fused_moe_lora_kernel.py &&
pytest -v -s lora/test_moe_lora_align_sum.py'
- label: LoRA Punica Kernels
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
set -o pipefail &&
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
- label: LoRA Punica FP8/XPU Ops
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s lora/test_punica_ops_fp8.py &&
pytest -v -s lora/test_punica_xpu_ops.py'
- label: LoRA Models
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
(pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
pytest -v -s lora/test_qwen35_densemodel_lora.py &&
pytest -v -s lora/test_transformers_model.py'
- label: LoRA Multimodal
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s lora/test_default_mm_loras.py &&
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
(pytest -v -s lora/test_qwenvl.py || true) &&
pytest -v -s lora/test_whisper.py'
-55
View File
@@ -1,55 +0,0 @@
group: Miscellaneous Intel
depends_on:
- image-build-xpu
steps:
- label: V1 Core + KV + Metrics
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/v1/core
- tests/v1/executor
- tests/v1/kv_offload
- tests/v1/worker
- tests/v1/kv_connector/unit
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install -r requirements/kv_connectors.txt &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
cd tests &&
pytest -v -s v1/executor'
- label: V1 Sample + Logits
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/v1/sample
- tests/v1/logits_processors
- tests/v1/test_oracle.py
- tests/v1/test_request.py
- tests/v1/test_outputs.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
cd tests &&
pytest -v -s v1/logits_processors &&
pytest -v -s v1/test_oracle.py &&
pytest -v -s v1/test_request.py &&
pytest -v -s v1/test_outputs.py'
+57 -62
View File
@@ -1,13 +1,3 @@
# CUDA architecture lists — following PyTorch RELEASE.md
# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
env:
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
# aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
steps:
- input: "Provide Release version here"
id: input-release-version
@@ -24,10 +14,12 @@ steps:
agents:
queue: arm64_cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
env:
DOCKER_BUILDKIT: "1"
@@ -37,7 +29,9 @@ steps:
agents:
queue: arm64_cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -63,7 +57,7 @@ steps:
agents:
queue: cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
@@ -76,7 +70,7 @@ steps:
agents:
queue: cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
@@ -114,95 +108,96 @@ steps:
depends_on: block-build-release-images
allow_dependency_failure: true
steps:
- label: "Build release image - x86_64 - CUDA 13.0"
- label: "Build release image - x86_64 - CUDA 12.9"
depends_on: ~
id: build-release-image-x86
agents:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- label: "Build release image - aarch64 - CUDA 13.0"
- label: "Build release image - aarch64 - CUDA 12.9"
depends_on: ~
id: build-release-image-arm64
agents:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
- label: "Build release image - x86_64 - CUDA 12.9"
- label: "Build release image - x86_64 - CUDA 13.0"
depends_on: ~
id: build-release-image-x86-cuda-12-9
id: build-release-image-x86-cuda-13-0
agents:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
- label: "Build release image - aarch64 - CUDA 12.9"
- label: "Build release image - aarch64 - CUDA 13.0"
depends_on: ~
id: build-release-image-arm64-cuda-12-9
id: build-release-image-arm64-cuda-13-0
agents:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-x86-ubuntu2404
agents:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-arm64-ubuntu2404
agents:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-x86-cuda-12-9-ubuntu2404
id: build-release-image-x86-cuda-13-0-ubuntu2404
agents:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-arm64-cuda-12-9-ubuntu2404
id: build-release-image-arm64-cuda-13-0-ubuntu2404
agents:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
- block: "Build release image for x86_64 CPU"
key: block-cpu-release-image-build
@@ -243,7 +238,7 @@ steps:
- group: "Publish release images"
key: "publish-release-images"
steps:
- label: "Create multi-arch manifest - CUDA 13.0"
- label: "Create multi-arch manifest - CUDA 12.9"
depends_on:
- build-release-image-x86
- build-release-image-arm64
@@ -255,7 +250,7 @@ steps:
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- label: "Annotate release workflow - CUDA 13.0"
- label: "Annotate release workflow - CUDA 12.9"
depends_on:
- create-multi-arch-manifest
id: annotate-release-workflow
@@ -264,19 +259,19 @@ steps:
commands:
- "bash .buildkite/scripts/annotate-release.sh"
- label: "Create multi-arch manifest - CUDA 12.9"
- label: "Create multi-arch manifest - CUDA 13.0"
depends_on:
- build-release-image-x86-cuda-12-9
- build-release-image-arm64-cuda-12-9
id: create-multi-arch-manifest-cuda-12-9
- build-release-image-x86-cuda-13-0
- build-release-image-arm64-cuda-13-0
id: create-multi-arch-manifest-cuda-13-0
agents:
queue: small_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
depends_on:
- build-release-image-x86-ubuntu2404
- build-release-image-arm64-ubuntu2404
@@ -288,17 +283,17 @@ steps:
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
depends_on:
- build-release-image-x86-cuda-12-9-ubuntu2404
- build-release-image-arm64-cuda-12-9-ubuntu2404
id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
- build-release-image-x86-cuda-13-0-ubuntu2404
- build-release-image-arm64-cuda-13-0-ubuntu2404
id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
agents:
queue: small_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
- label: "Publish nightly multi-arch image to DockerHub"
depends_on:
@@ -318,16 +313,16 @@ steps:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
depends_on:
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-cuda-13-0
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/push-nightly-builds.sh cu129"
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
# Clean up old nightly builds (keep only last 14)
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
plugins:
- docker-login#v3.0.0:
username: vllmbot
+22 -22
View File
@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
To download the wheel (by commit):
\`\`\`
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
(Optional) For CUDA 12.9:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
(Optional) For CUDA 13.0:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
(Optional) For CPU:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
@@ -33,8 +33,8 @@ To download and upload the image:
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai:latest-x86_64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker push vllm/vllm-openai:latest-x86_64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
docker push vllm/vllm-openai:latest-x86_64-cu130
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker push vllm/vllm-openai:latest-aarch64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
docker push vllm/vllm-openai:latest-aarch64-cu130
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
## ROCm
@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
docker manifest rm vllm/vllm-openai:latest-cu129
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest push vllm/vllm-openai:latest-cu129
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
docker manifest rm vllm/vllm-openai:latest-cu130
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
docker manifest push vllm/vllm-openai:latest-cu130
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
docker manifest rm vllm/vllm-openai-cpu:latest || true
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
TORCH_INDEX_URL=""
fi
else
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
fi
echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
@@ -9,7 +9,7 @@ set -ex
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
SUBPATH=$BUILDKITE_COMMIT
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
@@ -25,100 +25,22 @@ export PYTHONPATH=".."
###############################################################################
cleanup_docker() {
# Share the same lock with image pull to avoid cleanup/pull races on one node.
local docker_lock="/tmp/docker-pull.lock"
exec 9>"$docker_lock"
flock 9
docker_root=$(docker info -f '{{.DockerRootDir}}')
if [ -z "$docker_root" ]; then
echo "Failed to determine Docker root directory." >&2
flock -u 9
return 1
exit 1
fi
echo "Docker root directory: $docker_root"
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
threshold=70
if [ "$disk_usage" -gt "$threshold" ]; then
echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
docker image prune -f
docker volume prune -f && docker system prune --force --filter "until=72h" --all
echo "Docker images and volumes cleanup completed."
else
echo "Disk usage is below $threshold%. Checking old CI images anyway."
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
fi
echo "Old CI image cleanup completed."
flock -u 9
}
cleanup_old_ci_images() {
local repo_prefix="$1"
local current_image_ref="$2"
local ttl_hours="$3"
local aggressive_cleanup="$4"
if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
return 0
fi
if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
ttl_hours=72
fi
local now_epoch cutoff_epoch
now_epoch=$(date +%s)
cutoff_epoch=$((now_epoch - ttl_hours * 3600))
local -a used_image_ids
mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
local removed_count=0
local examined_count=0
declare -A seen_ids=()
while read -r image_ref image_id; do
[[ -z "$image_ref" || -z "$image_id" ]] && continue
((examined_count++))
# Keep the image this job is going to use.
if [[ "$image_ref" == "$current_image_ref" ]]; then
continue
fi
# Avoid duplicate deletes when multiple tags point to same image id.
if [[ -n "${seen_ids[$image_id]:-}" ]]; then
continue
fi
seen_ids[$image_id]=1
# Never delete images that are used by any container on this node.
if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
continue
fi
local created created_epoch
created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
[[ -z "$created" ]] && continue
created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
[[ -z "$created_epoch" ]] && continue
if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
if docker image rm -f "$image_id" >/dev/null 2>&1; then
((removed_count++))
fi
fi
done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
# Also trim old dangling layers; this is safe and does not remove referenced images.
docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
if [[ "$aggressive_cleanup" == "1" ]]; then
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
else
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
echo "Disk usage is below $threshold%. No cleanup needed."
fi
}
@@ -318,6 +240,7 @@ fi
cleanup_docker
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
# --- Build or pull test image ---
IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
@@ -343,6 +266,8 @@ fi
remove_docker_container() {
docker rm -f "${container_name}" || true
docker image rm -f "${image_name}" || true
docker system prune -f || true
}
trap remove_docker_container EXIT
@@ -358,7 +283,6 @@ docker run \
--ipc=host \
--privileged \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ${HOME}/.cache/huggingface:/root/.cache/huggingface \
--entrypoint="" \
-e "HF_TOKEN=${HF_TOKEN:-}" \
-e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
@@ -12,7 +12,9 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true
docker rm -f "${container_name}" || true;
docker image rm -f "${image_name}" || true;
docker system prune -f || true;
}
trap remove_docker_container EXIT
@@ -1,55 +0,0 @@
#!/usr/bin/env bash
set -euxo pipefail
# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
THRESHOLD=${1:-0.8}
NUM_Q=${2:-1319}
PORT=${3:-8050}
OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
mkdir -p "${OUT_DIR}"
wait_for_server() {
local port=$1
timeout 600 bash -c '
until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
sleep 1
done'
}
MODEL="Qwen/Qwen3-30B-A3B-FP8"
BACK="allgather_reducescatter"
cleanup() {
if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
kill "${SERVER_PID}" 2>/dev/null || true
for _ in {1..20}; do
kill -0 "${SERVER_PID}" 2>/dev/null || break
sleep 0.5
done
kill -9 "${SERVER_PID}" 2>/dev/null || true
fi
}
trap cleanup EXIT
VLLM_DEEP_GEMM_WARMUP=skip \
vllm serve "$MODEL" \
--enforce-eager \
--data-parallel-size 4 \
--enable-expert-parallel \
--enable-eplb \
--all2all-backend "$BACK" \
--eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \
--trust-remote-code \
--max-model-len 2048 \
--port "$PORT" &
SERVER_PID=$!
wait_for_server "$PORT"
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
python3 - <<PY
import json; acc=json.load(open('${OUT}'))['accuracy']
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
assert acc >= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}"
PY
+2924 -2571
View File
File diff suppressed because it is too large Load Diff
-98
View File
@@ -1,98 +0,0 @@
group: Disaggregated
depends_on:
- image-build
steps:
- label: Distributed NixlConnector PD accuracy (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
timeout_in_minutes: 30
device: a100
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+85
View File
@@ -226,6 +226,91 @@ steps:
commands:
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
- label: Distributed NixlConnector PD accuracy (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
timeout_in_minutes: 30
device: a100
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- label: Pipeline + Context Parallelism (4 GPUs)
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
@@ -29,15 +29,6 @@ steps:
commands:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
timeout_in_minutes: 60
device: h100
optional: true
num_devices: 4
working_dir: "/vllm-workspace"
commands:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
timeout_in_minutes: 60
device: h100
@@ -28,7 +28,6 @@ steps:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
mirror:
amd:
device: mi325_1
-34
View File
@@ -12,17 +12,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Eagle Nightly B200
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Speculators + MTP
timeout_in_minutes: 30
device: h200_18gb
@@ -34,18 +23,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Speculators + MTP Nightly B200
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/transformers_utils/configs/speculators/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Ngram + Suffix
timeout_in_minutes: 30
device: h200_18gb
@@ -66,17 +43,6 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Draft Model Nightly B200
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: DFlash Speculators Correctness
timeout_in_minutes: 30
device: h100
+1 -1
View File
@@ -262,7 +262,7 @@ pull_request_rules:
- files~=^docker/Dockerfile.xpu
- files~=^\\.buildkite/intel_jobs/
- files=\.buildkite/ci_config_intel.yaml
- files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
- files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
- files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
- files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
- files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+1 -1
View File
@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
# Limit the number of parallel jobs to avoid OOM
export MAX_JOBS=1
# Make sure release wheels are built for the following architectures
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
bash tools/check_repo.sh
+7 -13
View File
@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
# Supported python versions. These versions will be searched in order, the
# first match will be selected. These should be kept in sync with setup.py.
#
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
# Supported AMD GPU architectures.
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
# ROCm installation prefix. Default to /opt/rocm but allow override via
# -DROCM_PATH=/your/rocm/path when invoking cmake.
@@ -94,15 +94,12 @@ find_package(Torch REQUIRED)
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
# to support the whole generation without specifying all sub-architectures
# see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
else()
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
endif()
#
@@ -310,9 +307,7 @@ set(VLLM_EXT_SRC
"csrc/torch_bindings.cpp")
if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_EXT_SRC
"csrc/minimax_reduce_rms_kernel.cu"
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
@@ -1053,8 +1048,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_MOE_EXT_SRC
"csrc/moe/moe_wna16.cu"
"csrc/moe/grouped_topk_kernels.cu"
"csrc/moe/router_gemm.cu"
"csrc/moe/topk_softplus_sqrt_kernels.cu")
"csrc/moe/router_gemm.cu")
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
+2 -2
View File
@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
| <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
</p>
🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
---
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
- Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
@@ -404,7 +404,6 @@ def _build_attention_metadata(
query_start_loc=q_start_gpu,
query_start_loc_cpu=q_start_cpu,
seq_lens=seq_lens_gpu,
seq_lens_cpu_upper_bound=seq_lens_cpu,
_seq_lens_cpu=seq_lens_cpu,
_num_computed_tokens_cpu=num_computed_tokens_cpu,
slot_mapping=slot_mapping,
View File
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
from vllm.platforms import current_platform
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
fp8_w8a8_moe_quant_config,
nvfp4_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
CutlassExpertsFp4,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_experts,
fused_topk,
View File
-378
View File
@@ -1,378 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Generic benchmark harness for vLLM IR ops.
Usage:
python benchmarks/kernels/ir/bench_ir_ops.py
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul
python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph
python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/
"""
import argparse
import contextlib
import csv
import dataclasses
import datetime
import math
import os
import subprocess
import sys
import tempfile
# Ensure repo root is on sys.path so `benchmarks` is importable as a package.
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
if _REPO_ROOT not in sys.path:
sys.path.insert(0, _REPO_ROOT)
# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2
# directly by the dynamic linker, so Python-level sys.stderr redirect won't
# catch them).
_saved_fd = os.dup(2)
try:
with open(os.devnull, "w") as _devnull:
os.dup2(_devnull.fileno(), 2)
import torch
import vllm.kernels # noqa: E402, F401
finally:
os.dup2(_saved_fd, 2)
os.close(_saved_fd)
from tqdm import tqdm # noqa: E402
from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS # noqa: E402 # isort: skip
from vllm.ir.op import IrOp # noqa: E402
from vllm.platforms import current_platform # noqa: E402
from vllm.triton_utils import triton # noqa: E402
@dataclasses.dataclass(frozen=True)
class BenchConfig:
use_cuda_graph: bool = True
warmup: int = 25
rep: int = 100
def _pkg_version(name: str) -> str:
from importlib.metadata import PackageNotFoundError, version
with contextlib.suppress(PackageNotFoundError):
return version(name)
return "not installed"
_METADATA_LABELS = {
"timestamp": "Timestamp",
"git_commit": "Git commit",
"vllm": "vLLM",
"pytorch": "PyTorch",
"cuda_runtime": "CUDA runtime",
"triton": "Triton",
"cutlass": "CUTLASS",
"helion": "Helion",
"device": "Device",
"bench_mode": "Bench mode",
"warmup": "Warmup",
"rep": "Repetitions",
}
def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]:
from vllm.collect_env import get_env_info
env = get_env_info()
git_sha = "unknown"
with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError):
git_sha = (
subprocess.check_output(
["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
)
.decode()
.strip()
)
device_name = current_platform.get_device_name()
warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)"
rep_note = " replays" if cfg.use_cuda_graph else " ms"
return {
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"git_commit": git_sha,
"vllm": str(env.vllm_version),
"pytorch": str(env.torch_version),
"cuda_runtime": str(env.cuda_runtime_version),
"triton": triton.__version__,
"cutlass": _pkg_version("nvidia-cutlass-dsl"),
"helion": _pkg_version("helion"),
"device": device_name,
"bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager",
"warmup": f"{cfg.warmup}{warmup_note}",
"rep": f"{cfg.rep}{rep_note}",
}
def print_metadata(metadata: dict[str, str]):
print("=" * 60)
for key, val in metadata.items():
print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}")
print("=" * 60)
def _clone_args(args: tuple) -> tuple:
return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args)
# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as
# inplace=True will mutate bench_args across iterations. Both CUDA graph
# and eager modes will accumulate drift from repeated in-place mutation.
# We need to re-clone inputs per iteration for inplace ops.
def _bench_one(fn, args, cfg: BenchConfig) -> float:
bench_args = _clone_args(args)
bench_fn = lambda: fn(*bench_args)
if cfg.use_cuda_graph:
ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5])
else:
ms = triton.testing.do_bench(
bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5]
)
return ms * 1000
# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
# the benchmark baseline should be the compiled native (what vLLM runs by
# default) rather than the uncompiled native implementation.
def collect_timings(
op: IrOp, shape_configs: list[dict], cfg: BenchConfig
) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
def fmt(v) -> str:
return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v)
case_names = [
"_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
]
providers = [n for n, impl in op.impls.items() if impl.supported]
results: dict[str, dict[str, float]] = {c: {} for c in case_names}
for provider in providers:
impl = op.impls[provider]
desc = f"{op.name} / {provider}"
for case_name, kwargs in tqdm(
zip(case_names, shape_configs),
desc=desc,
total=len(case_names),
unit=" cases",
):
args = op.generate_inputs(**kwargs)
if impl.supports_args(*args):
results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
else:
results[case_name][provider] = float("nan")
return case_names, providers, results
def analyze_results(
op_name: str,
case_names: list[str],
providers: list[str],
results: dict[str, dict[str, float]],
) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
native_col = "native"
non_native = [p for p in providers if p != native_col]
header_cols = ["case"]
for p in providers:
header_cols.append(f"{p} (us)")
for p in non_native:
header_cols.append(f"{p} speedup")
detail_rows: list[dict[str, str]] = []
speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native}
for case_name in case_names:
timings = results[case_name]
row: dict[str, str] = {"case": case_name}
for p in providers:
val = timings.get(p, float("nan"))
row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a"
native_us = timings.get(native_col, float("nan"))
for p in non_native:
p_us = timings.get(p, float("nan"))
if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0:
speedup = native_us / p_us
row[f"{p} speedup"] = f"{speedup:.2f}x"
speedup_data[p].append((speedup, case_name))
else:
row[f"{p} speedup"] = "n/a"
detail_rows.append(row)
summary_rows: list[dict[str, str]] = []
for p in non_native:
entries = speedup_data[p]
if not entries:
continue
speedups = [s for s, _ in entries]
geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups))
best_val, best_case = max(entries)
worst_val, worst_case = min(entries)
wins = sum(1 for s in speedups if s > 1.0)
losses = sum(1 for s in speedups if s < 1.0)
total = len(speedups)
print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
print(f" geomean speedup: {geomean:.2f}x")
print(f" best: {best_val:.2f}x ({best_case})")
print(f" worst: {worst_val:.2f}x ({worst_case})")
summary_rows.append(
{
"op": op_name,
"provider": p,
"geomean_speedup": f"{geomean:.2f}",
"best_speedup": f"{best_val:.2f}",
"best_case": best_case,
"worst_speedup": f"{worst_val:.2f}",
"worst_case": worst_case,
"wins": str(wins),
"losses": str(losses),
"total": str(total),
}
)
return detail_rows, summary_rows, header_cols
def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]):
with open(path, "w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
def save_results(
save_dir: str,
op_name: str,
detail_rows: list[dict[str, str]],
header_cols: list[str],
all_summary_rows: list[dict[str, str]],
metadata: dict[str, str],
):
write_csv(
os.path.join(save_dir, f"{op_name}_detail.csv"),
detail_rows,
header_cols,
)
if all_summary_rows:
write_csv(
os.path.join(save_dir, "summary.csv"),
all_summary_rows,
list(all_summary_rows[0].keys()),
)
write_csv(
os.path.join(save_dir, "metadata.csv"),
[metadata],
list(metadata.keys()),
)
def parse_args():
parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
parser.add_argument(
"--ops",
type=str,
default=None,
help="Comma-separated list of op names to benchmark (substring match)",
)
parser.add_argument(
"--no-cuda-graph",
action="store_true",
help="Disable CUDA graph; use do_bench with L2 cache flushing instead",
)
parser.add_argument(
"--warmup",
type=int,
default=25,
help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)",
)
parser.add_argument(
"--rep",
type=int,
default=100,
help="Repetition time in ms (do_bench) or number of graph replays "
"(do_bench_cudagraph) (default: 100)",
)
parser.add_argument(
"--save-path",
type=str,
default=None,
help="Directory to save results (default: auto-created temp dir)",
)
return parser.parse_args()
def main():
args = parse_args()
cfg = BenchConfig(
use_cuda_graph=not args.no_cuda_graph,
warmup=args.warmup,
rep=args.rep,
)
torch.set_default_device(current_platform.device_type)
metadata = collect_env_metadata(cfg)
print_metadata(metadata)
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
save_dir = args.save_path or os.path.join(
tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}"
)
os.makedirs(save_dir, exist_ok=True)
op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
all_summary_rows: list[dict[str, str]] = []
for op in IrOp.registry.values():
if op_filters and not any(f in op.name for f in op_filters):
continue
if not op.has_input_generator:
print(f"Skipping op '{op.name}': no input generator registered")
continue
if op.name not in SHAPE_CONFIGS:
raise RuntimeError(
f"No benchmark shape config for op '{op.name}'. "
f"Add it to benchmarks/kernels/ir/shapes.py"
)
case_names, providers, results = collect_timings(
op, SHAPE_CONFIGS[op.name], cfg
)
detail_rows, summary_rows, header_cols = analyze_results(
op.name, case_names, providers, results
)
all_summary_rows.extend(summary_rows)
save_results(
save_dir,
op.name,
detail_rows,
header_cols,
all_summary_rows,
metadata,
)
print(f"\nResults saved to: {save_dir}")
if __name__ == "__main__":
main()
-29
View File
@@ -1,29 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Shape configurations for IR op benchmarks.
"""
import torch
NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384]
COMMON_HIDDEN_SIZES = [
2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n
3072, # Gemma 7B/9B
4096, # Llama 3 8B, Qwen 3 8B, Mistral 7B
5120, # Llama 4 Scout 17B-16E
7168, # DeepSeek V3
8192, # Llama 3 70B
16384, # Llama 3 405B
]
# Each entry maps an op name to a list of kwarg dicts that will be passed
# to that op's registered input generator via op.generate_inputs(**kwargs).
SHAPE_CONFIGS: dict[str, list[dict]] = {
"rms_norm": [
{"num_tokens": n, "hidden_size": d, "dtype": dtype}
for dtype in [torch.float16, torch.bfloat16, torch.float32]
for d in COMMON_HIDDEN_SIZES
for n in NUM_TOKENS
],
}
+8 -41
View File
@@ -161,49 +161,16 @@ elseif (S390_FOUND)
"-mtune=native")
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
message(STATUS "RISC-V detected")
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
# by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
if(NOT DEFINED VLLM_RVV_VLEN)
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
if(EXISTS /proc/cpuinfo)
file(READ /proc/cpuinfo _cpuinfo)
set(_best 0)
foreach(_n IN ITEMS 128 256 512 1024)
if(_cpuinfo MATCHES "zvl${_n}b")
set(_best ${_n})
endif()
endforeach()
if(_best GREATER 0)
set(VLLM_RVV_VLEN ${_best})
endif()
endif()
# If auto-detect failed (no /proc/cpuinfo or no zvl<N>b reported)
# but the compiler supports RVV, require explicit specification.
if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
message(FATAL_ERROR
"RISC-V RVV is available but VLEN could not be auto-detected. "
"Please specify VLEN explicitly:\n"
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n"
" -DVLLM_RVV_VLEN=0 (force scalar, no RVV)")
endif()
endif()
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
if(RVV_BF16_FOUND)
message(STATUS "BF16 extension detected")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
add_compile_definitions(RISCV_BF16_SUPPORT)
elseif(RVV_FP16_FOUND)
message(WARNING "BF16 functionality is not available")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
else()
message(STATUS "compile riscv with scalar (no FP16/BF16)")
set(MARCH_FLAGS -march=rv64gc)
endif()
if(RVV_BF16_FOUND)
message(STATUS "BF16 extension detected")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
add_compile_definitions(RISCV_BF16_SUPPORT)
elseif (RVV_FP16_FOUND)
message(WARNING "BF16 functionality is not available")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
else()
message(STATUS "compile riscv with scalar")
set(MARCH_FLAGS -march=rv64gc)
list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
endif()
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
else()
+1 -6
View File
@@ -20,7 +20,7 @@ else()
FetchContent_Declare(
deepgemm
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
@@ -120,11 +120,6 @@ if(DEEPGEMM_ARCHS)
COMPONENT _deep_gemm_C
FILES_MATCHING PATTERN "*.py")
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
DESTINATION vllm/third_party/deep_gemm/mega
COMPONENT _deep_gemm_C
FILES_MATCHING PATTERN "*.py")
# Generate envs.py (normally generated by DeepGEMM's setup.py build step)
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
"# Pre-installed environment variables\npersistent_envs = dict()\n")
+1 -1
View File
@@ -19,7 +19,7 @@ else()
FetchContent_Declare(
flashmla
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
+25 -82
View File
@@ -11,74 +11,29 @@
namespace vllm {
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y,
const float limit) {
if constexpr (act_first) {
scalar_t gate = x;
scalar_t up = y;
if constexpr (HAS_CLAMP) {
gate = (scalar_t)fminf((float)gate, limit);
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
}
return ACT_FN(gate) * up;
} else {
scalar_t gate = x;
scalar_t up = y;
if constexpr (HAS_CLAMP) {
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
up = (scalar_t)fminf((float)up, limit);
}
return gate * ACT_FN(up);
}
const scalar_t& y) {
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
}
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
const packed_t& y,
const float limit) {
if constexpr (act_first) {
packed_t gate = x;
packed_t up = y;
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fminf(g.x, limit);
g.y = fminf(g.y, limit);
u.x = fmaxf(fminf(u.x, limit), -limit);
u.y = fmaxf(fminf(u.y, limit), -limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(PACKED_ACT_FN(gate), up);
} else {
packed_t gate = x;
packed_t up = y;
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fmaxf(fminf(g.x, limit), -limit);
g.y = fmaxf(fminf(g.y, limit), -limit);
u.x = fminf(u.x, limit);
u.y = fminf(u.y, limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(gate, PACKED_ACT_FN(up));
}
const packed_t& y) {
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
: packed_mul(x, PACKED_ACT_FN(y));
}
// Activation and gating kernel template.
template <typename scalar_t, typename packed_t,
scalar_t (*ACT_FN)(const scalar_t&),
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
bool use_vec, bool use_256b = false>
__global__ void act_and_mul_kernel(
scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d, const float limit) {
const int d) {
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
}
#pragma unroll
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
x.elts[j] =
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
x.elts[j], y.elts[j], limit);
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
x.elts[j], y.elts[j]);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
out_ptr[idx] =
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
}
}
}
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
// Launch activation and gating kernel.
// Use ACT_FIRST (bool) indicating whether to apply the activation function
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
// clamped (max only) and up input is clamped (both sides) before the
// activation function is applied.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
HAS_CLAMP, LIMIT) \
// first.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
auto dtype = input.scalar_type(); \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} \
} else { \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
}
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, false, 0.0f);
}
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
double limit) {
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, true, (float)limit);
true);
}
void mul_and_silu(torch::Tensor& out, // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
// applies the silu to the latter half of the input.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
false, false, 0.0f);
false);
}
void gelu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
true, false, 0.0f);
true);
}
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true);
}
namespace vllm {
+8 -18
View File
@@ -599,11 +599,6 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
// Find batch index within a block
__shared__ int batch_idx[BLOCK_Y_SIZE];
if (threadIdx.x == 0) {
batch_idx[threadIdx.y] = -1;
}
__syncthreads();
for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
iter++) {
int tid = iter * blockDim.x + threadIdx.x;
@@ -616,18 +611,16 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
}
}
__syncthreads();
#ifndef USE_ROCM
__syncwarp();
#endif
// num_tokens may be an allocation upper bound when Python avoids a D2H sync.
// Only tokens covered by the exact device-side cu_seq_lens are valid to
// gather.
const int batch = batch_idx[threadIdx.y];
if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
if (head_idx >= head_dim || token_idx >= num_tokens) {
return;
}
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
const int block_idx =
block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
inbatch_seq_idx / cache_block_size];
const int64_t src_block_offset = block_idx * block_stride;
const int64_t cache_inblock_offset =
(inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
@@ -1497,9 +1490,6 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
ql_nope.scalar_type() == at::ScalarType::BFloat16,
"ql_nope must be float16 or bfloat16 dtype");
if (num_tokens == 0) return;
@@ -1511,7 +1501,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
+2 -48
View File
@@ -61,23 +61,8 @@
#endif
#ifdef __aarch64__
// Implementation of neon_expf copied from Arm Optimized Routines (expf
// AdvSIMD)
// Implementation copied from Arm Optimized Routines (expf AdvSIMD)
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
//
// Additional fast exponential intended for cases where outputs will be
// downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
// for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
// clamps inputs outside this range to 0 / inf. Implementation is similar to
// exp_u20, but:
// - uses a third degree polynomial approximation for exp(r) instead of a
// fifth degree one, with coefficients re-tuned.
// - does not split natural log (ln) into high / low parts
// - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
// exp(x) = 2^n (exp(r))
// r = x - n*ln2, with n = round(x/ln2)
// exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
// n = round(x / ln2), r = x - n*ln2
#include <limits>
#define DEFINE_FAST_EXP \
const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
@@ -121,38 +106,7 @@
result.val[2] = neon_expf(vec.reg.val[2]); \
result.val[3] = neon_expf(vec.reg.val[3]); \
return vec_op::FP32Vec16(result); \
}; \
const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
constexpr float ln2 = 0x1.62e43p-1f; \
constexpr float f_c2 = 0x1.5592ecp-3f; \
const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
always_inline)) { \
const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
float32x4_t r = vfmsq_n_f32(values, n, ln2); \
uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
float32x4_t r2 = vmulq_f32(r, r); \
float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
float32x4_t p = vfmaq_f32(s, q, r2); \
float32x4_t y = \
vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
return y; \
}; \
auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
__attribute__((always_inline)) { \
float32x4x4_t result; \
result.val[0] = neon_expf_f16(vec.reg.val[0]); \
result.val[1] = neon_expf_f16(vec.reg.val[1]); \
result.val[2] = neon_expf_f16(vec.reg.val[2]); \
result.val[3] = neon_expf_f16(vec.reg.val[3]); \
return vec_op::FP32Vec16(result); \
};
};
#endif // __aarch64__
+4 -28
View File
@@ -1152,11 +1152,7 @@ class AttentionMainLoop {
bool use_sink) {
#ifdef DEFINE_FAST_EXP
DEFINE_FAST_EXP
bool constexpr IsReducedPrecision =
std::is_same_v<query_t, c10::BFloat16> ||
std::is_same_v<query_t, c10::Half>;
#endif
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
@@ -1205,17 +1201,8 @@ class AttentionMainLoop {
vec = vec - max_vec;
// compute exp
#if defined(DEFINE_FAST_EXP)
#ifdef __aarch64__
if constexpr (IsReducedPrecision) {
vec = fast_exp_f16(vec);
} else
#endif
{
vec = fast_exp(vec);
}
#ifdef DEFINE_FAST_EXP
vec = fast_exp(vec);
prob_buffer_vec_t output_vec(vec);
output_vec.save(curr_prob_buffer_iter);
#else
@@ -1271,11 +1258,7 @@ class AttentionMainLoop {
int32_t kv_tile_token_num, float softcap_scale) {
#ifdef DEFINE_FAST_EXP
DEFINE_FAST_EXP
bool constexpr IsReducedPrecision =
std::is_same_v<query_t, c10::BFloat16> ||
std::is_same_v<query_t, c10::Half>;
#endif
float inv_softcap_scale = 1.0 / softcap_scale;
vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
@@ -1289,15 +1272,8 @@ class AttentionMainLoop {
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
vec = vec * inv_softcap_scale_vec;
#if defined(DEFINE_FAST_EXP)
#ifdef __aarch64__
if constexpr (IsReducedPrecision) {
vec = fast_exp_f16(vec);
} else
#endif
{
vec = fast_exp(vec);
}
#ifdef DEFINE_FAST_EXP
vec = fast_exp(vec);
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
vec = (vec - inv_vec) / (vec + inv_vec);
#else
+823 -16
View File
@@ -1,25 +1,832 @@
#ifndef CPU_TYPES_RISCV_HPP
#define CPU_TYPES_RISCV_HPP
// RISC-V Vector (RVV) CPU type definitions for vLLM.
//
// Supports multiple VLENs via compile-time dispatch. The compiler defines
// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
// maps VLEN to the correct LMUL suffixes, and the impl header provides
// VLEN-independent class implementations.
//
// To add support for a new VLEN, add the LMUL mapping in
// cpu_types_riscv_defs.hpp (the impl header needs no changes).
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iostream>
#include <limits>
#include <riscv_vector.h>
#include <torch/all.h>
#ifndef __riscv_vector
#error "cpu_types_riscv.hpp included in a non-RVV translation unit"
// ============================================================================
// Vector Register Type Definitions (VLEN=128 bits)
// ============================================================================
typedef vfloat16m1_t fixed_vfloat16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vfloat16m2_t fixed_vfloat16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vfloat32m1_t fixed_vfloat32m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vfloat32m2_t fixed_vfloat32m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vfloat32m4_t fixed_vfloat32m4_t
__attribute__((riscv_rvv_vector_bits(512)));
typedef vfloat32m8_t fixed_vfloat32m8_t
__attribute__((riscv_rvv_vector_bits(1024)));
typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
typedef vuint16m1_t fixed_vuint16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vuint16m2_t fixed_vuint16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vuint16m4_t fixed_vuint16m4_t
__attribute__((riscv_rvv_vector_bits(512)));
#ifdef RISCV_BF16_SUPPORT
typedef vbfloat16m1_t fixed_vbfloat16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vbfloat16m2_t fixed_vbfloat16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vbfloat16m4_t fixed_vbfloat16m4_t
__attribute__((riscv_rvv_vector_bits(512)));
#endif
#ifndef __riscv_v_min_vlen
#error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
namespace vec_op {
#ifdef RISCV_BF16_SUPPORT
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#else
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
#endif
#include "cpu_types_riscv_defs.hpp"
#include "cpu_types_riscv_impl.hpp"
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#endif // CPU_TYPES_RISCV_HPP
#define FORCE_INLINE __attribute__((always_inline)) inline
namespace {
template <typename T, T... indexes, typename F>
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
(f(std::integral_constant<T, indexes>{}), ...);
};
} // namespace
template <typename T, T count, typename F,
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
constexpr void unroll_loop(F&& f) {
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
}
template <typename T>
struct Vec {
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
};
struct FP32Vec8;
struct FP32Vec16;
// ============================================================================
// FP16 Implementation
// ============================================================================
struct FP16Vec8 : public Vec<FP16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat16m1_t reg;
explicit FP16Vec8(const void* ptr)
: reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
VEC_ELEM_NUM)) {};
explicit FP16Vec8(const FP32Vec8&);
void save(void* ptr) const {
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
__riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
VEC_ELEM_NUM);
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat16m2_t reg;
explicit FP16Vec16(const void* ptr)
: reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
VEC_ELEM_NUM)) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
__riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
VEC_ELEM_NUM);
}
};
// ============================================================================
// BF16 Implementation
// ============================================================================
#ifdef RISCV_BF16_SUPPORT
FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
return __riscv_vreinterpret_v_bf16m1_u16m1(v);
}
FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
return __riscv_vreinterpret_v_bf16m2_u16m2(v);
}
FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
return __riscv_vreinterpret_v_bf16m4_u16m4(v);
}
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vbfloat16m1_t reg;
explicit BF16Vec8(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vbfloat16m2_t reg;
explicit BF16Vec16(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_vbfloat16m4_t reg;
explicit BF16Vec32(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
explicit BF16Vec32(const BF16Vec8& v) {
fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
fixed_vuint16m4_t u16_combined =
__riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
};
void save(void* ptr) const {
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
#else
// ============================================================================
// BF16 Fallback Implementation (FP32 Simulation)
// ============================================================================
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat32m2_t reg_fp32;
explicit BF16Vec8(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[8];
for (int i = 0; i < 8; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
}
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat32m4_t reg_fp32;
explicit BF16Vec16(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[16];
for (int i = 0; i < 16; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_vfloat32m8_t reg_fp32;
explicit BF16Vec32(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[32];
for (int i = 0; i < 32; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
}
explicit BF16Vec32(const BF16Vec8& v) {
float tmp_small[8];
__riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
float tmp_large[32];
for (int i = 0; i < 4; ++i) {
std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
}
reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
}
void save(void* ptr) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
#endif
// ============================================================================
// FP32 Implementation
// ============================================================================
struct FP32Vec4 : public Vec<FP32Vec4> {
constexpr static int VEC_ELEM_NUM = 4;
fixed_vfloat32m1_t reg;
explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec4(const float* ptr)
: reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m1(ptr, reg, elem_num);
}
};
struct FP32Vec8 : public Vec<FP32Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat32m2_t reg;
explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const float* ptr)
: reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
explicit FP32Vec8(const FP16Vec8& v)
: reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_vfloat16m1_t v)
: reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec8(fixed_vbfloat16m1_t v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const BF16Vec8& v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
#else
explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
#endif
float reduce_sum() const {
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
FP32Vec8 operator*(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator+(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator-(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator/(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 max(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 abs() const {
return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
}
FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
}
FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
fixed_vfloat32m2_t temp =
__riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
}
void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m2(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
__riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec8 exp() const {
const float inv_ln2 = 1.44269504088896341f;
fixed_vfloat32m2_t x_scaled =
__riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
fixed_vfloat32m2_t r =
__riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
fixed_vfloat32m2_t poly =
__riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
fixed_vint32m2_t biased_exp =
__riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
fixed_vint32m2_t exponent_bits =
__riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
fixed_vfloat32m2_t scale =
__riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
}
FP32Vec8 tanh() const {
fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
__riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t x2 =
__riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
FP32Vec8 exp_val = FP32Vec8(x2).exp();
fixed_vfloat32m2_t num =
__riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t den =
__riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
}
FP32Vec8 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
__riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a4, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a3, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a2, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a1, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
fixed_vfloat32m2_t exp_val =
FP32Vec8(__riscv_vfneg_v_f32m2(
__riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
__riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
}
};
struct FP32Vec16 : public Vec<FP32Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat32m4_t reg;
explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const float* ptr)
: reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
explicit FP32Vec16(const FP32Vec8& data)
: reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(const FP16Vec16& v);
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec16(fixed_vbfloat16m2_t v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const BF16Vec16& v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
#else
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator-(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator*(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator/(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
}
float reduce_sum() const {
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_max() const {
fixed_vfloat32m1_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_min() const {
fixed_vfloat32m1_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
template <int group_size>
float reduce_sub_sum(int idx) {
static_assert(VEC_ELEM_NUM % group_size == 0);
const int start = idx * group_size;
vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
vbool8_t mask = __riscv_vmand_mm_b8(
__riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
__riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
VEC_ELEM_NUM);
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar =
__riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
};
FP32Vec16 max(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 abs() const {
return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
}
FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
return FP32Vec16(__riscv_vfmin_vv_f32m4(
max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
VEC_ELEM_NUM));
}
void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m4(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
__riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec16 exp() const {
const float inv_ln2 = 1.44269504088896341f;
fixed_vfloat32m4_t x_scaled =
__riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
fixed_vfloat32m4_t r =
__riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
fixed_vfloat32m4_t poly =
__riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.009618129107628f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.055504108664821f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.240226506959101f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.693147180559945f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM);
fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
__riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
__riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
}
FP32Vec16 tanh() const {
fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
__riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
FP32Vec16 exp_val =
FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
return FP32Vec16(__riscv_vfdiv_vv_f32m4(
__riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
__riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
}
FP32Vec16 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
__riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM);
fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a4, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a3, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a2, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a1, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
fixed_vfloat32m4_t exp_val =
FP32Vec16(__riscv_vfneg_v_f32m4(
__riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
__riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
}
};
// ============================================================================
// Type Traits & Global Helpers
// ============================================================================
template <typename T>
struct VecType {
using vec_type = void;
using vec_t = void;
};
template <typename T>
using vec_t = typename VecType<T>::vec_type;
template <>
struct VecType<float> {
using vec_type = FP32Vec8;
using vec_t = FP32Vec8;
};
template <>
struct VecType<c10::Half> {
using vec_type = FP16Vec8;
using vec_t = FP16Vec8;
};
template <>
struct VecType<c10::BFloat16> {
using vec_type = BF16Vec8;
using vec_t = BF16Vec8;
};
template <typename T>
void storeFP32(float v, T* ptr) {
*ptr = v;
}
template <>
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
*reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
}
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
}
inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
#ifdef RISCV_BF16_SUPPORT
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
*ptr = static_cast<__bf16>(v);
};
inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
: reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
: reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
#else
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
uint32_t val;
std::memcpy(&val, &v, 4);
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
#endif
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
} // namespace vec_op
#ifndef CPU_KERNEL_GUARD_IN
#define CPU_KERNEL_GUARD_IN(NAME)
#endif
#ifndef CPU_KERNEL_GUARD_OUT
#define CPU_KERNEL_GUARD_OUT(NAME)
#endif
#endif // CPU_TYPES_RISCV_HPP
-98
View File
@@ -1,98 +0,0 @@
#ifndef CPU_TYPES_RISCV_DEFS_HPP
#define CPU_TYPES_RISCV_DEFS_HPP
// VLEN-to-LMUL mapping for RISC-V Vector extension.
//
// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
#include <riscv_vector.h>
#if __riscv_v_min_vlen == 128
#define LMUL_128 m1
#define LMUL_256 m2
#define LMUL_512 m4
#define LMUL_1024 m8
#define BOOL_256 b16
#define BOOL_512 b8
#elif __riscv_v_min_vlen == 256
#define LMUL_128 mf2
#define LMUL_256 m1
#define LMUL_512 m2
#define LMUL_1024 m4
#define BOOL_256 b32
#define BOOL_512 b16
#else
#error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
#endif
// Token-paste helpers.
#define _RVV_P2(a, b) a##b
#define _RVV_P3(a, b, c) a##b##c
#define _RVV_P4(a, b, c, d) a##b##c##d
#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
#define RVVI(base, lmul) _RVV_P2(base, lmul)
#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
// ---- Semantic fixed-vector typedefs (named by element count) ----
// float16
typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
__attribute__((riscv_rvv_vector_bits(256)));
// float32
typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
__attribute__((riscv_rvv_vector_bits(512)));
typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
__attribute__((riscv_rvv_vector_bits(1024)));
// int32
typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
__attribute__((riscv_rvv_vector_bits(512)));
// uint16
typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
__attribute__((riscv_rvv_vector_bits(512)));
// bfloat16
#ifdef RISCV_BF16_SUPPORT
typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
__attribute__((riscv_rvv_vector_bits(512)));
#endif
// ---- Reduction accumulator type (always m1 = one register of f32) ----
// Used for scalar reductions; only element [0] is meaningful.
typedef vfloat32m1_t rvv_f32_accum_t
__attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
// ---- Mask types for f32 elements ----
#if __riscv_v_min_vlen == 128
typedef vbool16_t rvv_mask_f32x8_t;
typedef vbool8_t rvv_mask_f32x16_t;
#elif __riscv_v_min_vlen == 256
typedef vbool32_t rvv_mask_f32x8_t;
typedef vbool16_t rvv_mask_f32x16_t;
#endif
#endif // CPU_TYPES_RISCV_DEFS_HPP
-905
View File
@@ -1,905 +0,0 @@
#ifndef CPU_TYPES_RISCV_IMPL_HPP
#define CPU_TYPES_RISCV_IMPL_HPP
// Shared implementation of RVV vector-type wrapper classes.
// This file is VLEN-independent: it uses the semantic type names and
// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp.
//
// DO NOT include this file directly; include cpu_types_riscv.hpp instead.
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iostream>
#include <limits>
#include <torch/all.h>
namespace vec_op {
// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
// is defined, otherwise via the FP32-simulation fallback path.
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#define FORCE_INLINE __attribute__((always_inline)) inline
namespace {
template <typename T, T... indexes, typename F>
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
(f(std::integral_constant<T, indexes>{}), ...);
};
} // namespace
template <typename T, T count, typename F,
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
constexpr void unroll_loop(F&& f) {
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
}
template <typename T>
struct Vec {
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
};
struct FP32Vec8;
struct FP32Vec16;
// ============================================================================
// FP16 Implementation
// ============================================================================
struct FP16Vec8 : public Vec<FP16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_fp16x8_t reg;
explicit FP16Vec8(const void* ptr)
: reg(RVVI(__riscv_vle16_v_f16, LMUL_128)(
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
explicit FP16Vec8(const FP32Vec8&);
void save(void* ptr) const {
RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr),
byte_stride, reg, VEC_ELEM_NUM);
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_fp16x16_t reg;
explicit FP16Vec16(const void* ptr)
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr),
byte_stride, reg, VEC_ELEM_NUM);
}
};
// ============================================================================
// BF16 Implementation
// ============================================================================
#ifdef RISCV_BF16_SUPPORT
FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
}
FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) {
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v);
}
FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) {
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v);
}
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_bf16x8_t reg;
explicit BF16Vec8(const void* ptr)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16,
LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {};
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
byte_stride, bf16_to_u16(reg),
VEC_ELEM_NUM);
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_bf16x16_t reg;
explicit BF16Vec16(const void* ptr)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16,
LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
byte_stride, bf16_to_u16(reg),
VEC_ELEM_NUM);
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_bf16x32_t reg;
explicit BF16Vec32(const void* ptr)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
explicit BF16Vec32(const BF16Vec8& v) {
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
fixed_u16x32_t u16_combined =
RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
u16_val, u16_val, u16_val, u16_val);
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
LMUL_512)(u16_combined);
};
void save(void* ptr) const {
RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
bf16_to_u16(reg), elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
byte_stride, bf16_to_u16(reg),
VEC_ELEM_NUM);
}
};
#else
// ============================================================================
// BF16 Fallback Implementation (FP32 Simulation)
// ============================================================================
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_fp32x8_t reg_fp32;
explicit BF16Vec8(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[8];
for (int i = 0; i < 8; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
}
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
float tmp[8];
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[8];
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[8];
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_fp32x16_t reg_fp32;
explicit BF16Vec16(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[16];
for (int i = 0; i < 16; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[16];
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[16];
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_fp32x32_t reg_fp32;
explicit BF16Vec32(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[32];
for (int i = 0; i < 32; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
}
explicit BF16Vec32(const BF16Vec8& v) {
float tmp_small[8];
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
float tmp_large[32];
for (int i = 0; i < 4; ++i) {
std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32);
}
void save(void* ptr) const {
float tmp[32];
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[32];
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[32];
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
#endif
// ============================================================================
// FP32 Implementation
// ============================================================================
struct FP32Vec4 : public Vec<FP32Vec4> {
constexpr static int VEC_ELEM_NUM = 4;
fixed_fp32x4_t reg;
explicit FP32Vec4(float v)
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {};
explicit FP32Vec4()
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec4(const float* ptr)
: reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {};
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
void save(float* ptr) const {
RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
}
void save(float* ptr, int elem_num) const {
RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num);
}
};
struct FP32Vec8 : public Vec<FP32Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_fp32x8_t reg;
explicit FP32Vec8(float v)
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8()
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const float* ptr)
: reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {};
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
explicit FP32Vec8(const FP16Vec8& v)
: reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_fp16x8_t v)
: reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec8(fixed_bf16x8_t v)
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const BF16Vec8& v)
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
};
#else
explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
#endif
float reduce_sum() const {
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar,
VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
FP32Vec8 operator*(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator+(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator-(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator/(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 max(const FP32Vec8& b) const {
return FP32Vec8(
RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 abs() const {
return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num));
}
FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num));
}
FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
fixed_fp32x8_t temp =
RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM);
return FP32Vec8(
RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM));
}
void save(float* ptr) const {
RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM);
}
void save(float* ptr, int elem_num) const {
RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec8 exp() const {
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
VEC_ELEM_NUM);
const float inv_ln2 = 1.44269504088896341f;
fixed_fp32x8_t x_scaled =
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
fixed_i32x8_t n_int =
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
fixed_fp32x8_t n_float =
RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM);
fixed_fp32x8_t r =
RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM);
fixed_fp32x8_t poly =
RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM);
fixed_i32x8_t biased_exp =
RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM);
biased_exp =
RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM);
fixed_i32x8_t exponent_bits =
RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM);
fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32,
LMUL_256)(exponent_bits);
return FP32Vec8(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM));
}
FP32Vec8 tanh() const {
fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
VEC_ELEM_NUM);
fixed_fp32x8_t x2 =
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM);
FP32Vec8 exp_val = FP32Vec8(x2).exp();
fixed_fp32x8_t num =
RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
fixed_fp32x8_t den =
RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
return FP32Vec8(
RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM));
}
FP32Vec8 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_fp32x8_t abs_x =
RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM);
fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM);
fixed_fp32x8_t poly =
RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM);
fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(
abs_x, abs_x, VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)(
RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)(
reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec8(
RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM));
}
};
struct FP32Vec16 : public Vec<FP32Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_fp32x16_t reg;
explicit FP32Vec16(float v)
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16()
: reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const float* ptr)
: reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {};
explicit FP32Vec16(const FP32Vec8& data)
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(const FP16Vec16& v);
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec16(fixed_bf16x16_t v)
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const BF16Vec16& v)
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) {
};
#else
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator-(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator*(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator/(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM));
}
float reduce_sum() const {
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar,
VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_max() const {
rvv_f32_accum_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar,
VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_min() const {
rvv_f32_accum_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar,
VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
template <int group_size>
float reduce_sub_sum(int idx) {
static_assert(VEC_ELEM_NUM % group_size == 0);
const int start = idx * group_size;
auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM);
rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)(
RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start,
VEC_ELEM_NUM),
RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)(
indices, start + group_size, VEC_ELEM_NUM),
VEC_ELEM_NUM);
rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)(
mask, reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
};
FP32Vec16 max(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 abs() const {
return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
}
FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)(
max_v.reg,
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM),
VEC_ELEM_NUM));
}
void save(float* ptr) const {
RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM);
}
void save(float* ptr, int elem_num) const {
RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec16 exp() const {
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
VEC_ELEM_NUM);
const float inv_ln2 = 1.44269504088896341f;
fixed_fp32x16_t x_scaled =
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
fixed_i32x16_t n_int =
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
fixed_fp32x16_t n_float =
RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM);
fixed_fp32x16_t r =
RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM);
fixed_fp32x16_t poly =
RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
0.009618129107628f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
0.055504108664821f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
0.240226506959101f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
0.693147180559945f, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)(
RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0,
VEC_ELEM_NUM);
fixed_fp32x16_t scale =
RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)(
RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM));
return FP32Vec16(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM));
}
FP32Vec16 tanh() const {
fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
VEC_ELEM_NUM);
FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)(
x_clamped, 2.0f, VEC_ELEM_NUM))
.exp();
return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(
RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
VEC_ELEM_NUM));
}
FP32Vec16 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_fp32x16_t abs_x =
RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM);
fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)(
RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM);
fixed_fp32x16_t poly =
RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1,
VEC_ELEM_NUM);
poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM);
fixed_fp32x16_t exp_val =
FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x,
VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)(
RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)(
reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec16(
RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM));
}
};
// ============================================================================
// Type Traits & Global Helpers
// ============================================================================
template <typename T>
struct VecType {
using vec_type = void;
using vec_t = void;
};
template <typename T>
using vec_t = typename VecType<T>::vec_type;
template <>
struct VecType<float> {
using vec_type = FP32Vec8;
using vec_t = FP32Vec8;
};
template <>
struct VecType<c10::Half> {
using vec_type = FP16Vec8;
using vec_t = FP16Vec8;
};
template <>
struct VecType<c10::BFloat16> {
using vec_type = BF16Vec8;
using vec_t = BF16Vec8;
};
template <typename T>
void storeFP32(float v, T* ptr) {
*ptr = v;
}
template <>
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
*reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM);
}
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM);
}
inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
#ifdef RISCV_BF16_SUPPORT
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
*ptr = static_cast<__bf16>(v);
};
inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
: reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) {
};
inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
: reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
};
#else
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
uint32_t val;
std::memcpy(&val, &v, 4);
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
#endif
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
} // namespace vec_op
#ifndef CPU_KERNEL_GUARD_IN
#define CPU_KERNEL_GUARD_IN(NAME)
#endif
#ifndef CPU_KERNEL_GUARD_OUT
#define CPU_KERNEL_GUARD_OUT(NAME)
#endif
#endif // CPU_TYPES_RISCV_IMPL_HPP
+1 -6
View File
@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
std::optional<torch::Tensor> key, int64_t head_size,
torch::Tensor& cos_sin_cache, bool is_neox,
int64_t rope_dim_offset, bool inverse) {
TORCH_CHECK(rope_dim_offset == 0,
"rope_dim_offset != 0 is not supported on CPU");
TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
torch::Tensor& cos_sin_cache, bool is_neox) {
int num_tokens = positions.numel();
int rot_dim = cos_sin_cache.size(1);
int num_heads = query.size(-1) / head_size;
+1 -2
View File
@@ -263,8 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
" Tensor cos_sin_cache, bool is_neox, int "
"rope_dim_offset=0, bool inverse=False) -> ()");
" Tensor cos_sin_cache, bool is_neox) -> ()");
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
// Quantization
@@ -1,477 +0,0 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* Horizontally-fused DeepseekV4-MLA kernel:
* - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
* - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
* cache insert
*
* Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
* TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
* dispatch choosing Q vs KV work per warp. The per-warp RMSNorm/RoPE
* skeleton is adapted from vllm-deepseek_v4's existing
* `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
*
* Assumptions (hard-coded for DeepseekV4 attention):
* HEAD_DIM = 512
* ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
* NOPE_DIM = 448
* QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
* FP8_MAX = 448.0f
* is_neox=false (GPT-J interleaved pairs)
* 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
*/
#include <cmath>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <type_traits>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/cuda.h>
#include "cuda_compat.h"
#include "dispatch_utils.h"
#include "type_convert.cuh"
#ifndef FINAL_MASK
#define FINAL_MASK 0xffffffffu
#endif
namespace vllm {
namespace deepseek_v4_fused_ops {
namespace {
inline int getSMVersion() {
auto* props = at::cuda::getCurrentDeviceProperties();
return props->major * 10 + props->minor;
}
} // namespace
// ────────────────────────────────────────────────────────────────────────────
// Constants
// ────────────────────────────────────────────────────────────────────────────
constexpr int kHeadDim = 512;
constexpr int kRopeDim = 64;
constexpr int kNopeDim = kHeadDim - kRopeDim; // 448
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
constexpr float kFp8Max = 448.0f;
// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
constexpr int kNumLanes = 32;
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
// ────────────────────────────────────────────────────────────────────────────
// Small inline helpers
// ────────────────────────────────────────────────────────────────────────────
__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);
val = fmaxf(val, peer);
peer = __shfl_xor_sync(FINAL_MASK, val, 2);
val = fmaxf(val, peer);
return val;
}
template <typename T>
__device__ __forceinline__ float warpSum(float val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
}
return val;
}
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
//
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
// warp handles one (token, head_slot) pair. head_slot < num_heads_q →
// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV
// branch (RoPE + UE8M0 quant + insert)
//
// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
// The Q branch covers all `num_tokens_full` rows (downstream attention uses
// them). The KV branch only inserts the first `num_tokens_insert` tokens
// (= slot_mapping length) into the paged cache.
//
template <typename scalar_t_in>
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
uint8_t* __restrict__ k_cache, // [num_blocks, block_stride]
int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
int64_t const* __restrict__ position_ids, // [N] i64
float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
float const eps,
int const num_tokens_full, // = q.size(0) = kv.size(0)
int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
int const num_heads_q, // H
int const cache_block_size, // tokens per paged-cache block
int const kv_block_stride) { // bytes per paged-cache block
#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
// no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
return;
} else {
#endif
using Converter = vllm::_typeConvert<scalar_t_in>;
int const warpsPerBlock = blockDim.x / 32;
int const warpId = threadIdx.x / 32;
int const laneId = threadIdx.x % 32;
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
int const total_slots_per_token = num_heads_q + 1;
int const tokenIdx = globalWarpIdx / total_slots_per_token;
int const slotIdx = globalWarpIdx % total_slots_per_token;
if (tokenIdx >= num_tokens_full) return;
bool const isKV = (slotIdx == num_heads_q);
// KV branch: skip DP-padded tokens (no slot reserved for them).
if (isKV && tokenIdx >= num_tokens_insert) return;
// PDL: wait for predecessor kernel (upstream q/kv producer) to signal
// before touching any global memory. No-op when PDL is not enabled on
// the launch. The CUDA runtime wrapper emits the griddepcontrol.wait
// PTX with the required memory clobber internally.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
// Dim range this lane owns within the 512-wide head.
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
// ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
float elements[kElemsPerLane];
float sumOfSquares = 0.0f;
scalar_t_in const* src_ptr;
if (isKV) {
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
} else {
int64_t const q_row_offset =
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
dim_base;
src_ptr = q_inout + q_row_offset;
}
// Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
// type and bitcast to scalar_t_in packed pairs for conversion.
uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
{
typename Converter::packed_hip_type const* p0 =
reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
typename Converter::packed_hip_type const* p1 =
reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
#pragma unroll
for (int i = 0; i < 4; i++) {
float2 f2 = Converter::convert(p0[i]);
elements[2 * i] = f2.x;
elements[2 * i + 1] = f2.y;
}
#pragma unroll
for (int i = 0; i < 4; i++) {
float2 f2 = Converter::convert(p1[i]);
elements[8 + 2 * i] = f2.x;
elements[8 + 2 * i + 1] = f2.y;
}
}
// ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
// Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
// The downstream bf16 round only happens at the final store.
if (!isKV) {
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
sumOfSquares += elements[i] * elements[i];
}
sumOfSquares = warpSum<float>(sumOfSquares);
float const rms_rcp =
rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
elements[i] = elements[i] * rms_rcp;
}
}
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
// All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
bool const is_rope_lane = dim_base >= kNopeDim;
if (is_rope_lane) {
int64_t const pos = position_ids[tokenIdx];
constexpr int kHalfRope = kRopeDim / 2; // 32
float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
float const* sin_ptr = cos_ptr + kHalfRope;
int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16
#pragma unroll
for (int p = 0; p < kElemsPerLane / 2; p++) {
int const pair_dim = rope_local_base + 2 * p;
int const half_idx = pair_dim / 2;
float const cos_v = VLLM_LDG(cos_ptr + half_idx);
float const sin_v = VLLM_LDG(sin_ptr + half_idx);
float const x_even = elements[2 * p];
float const x_odd = elements[2 * p + 1];
elements[2 * p] = x_even * cos_v - x_odd * sin_v;
elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
}
}
// ═══════════════════════════════════════════════════════════════════════
// Q branch: cast to bf16 and store back in place.
// ═══════════════════════════════════════════════════════════════════════
if (!isKV) {
uint4 out0, out1;
typename Converter::packed_hip_type* po0 =
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
typename Converter::packed_hip_type* po1 =
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
#pragma unroll
for (int i = 0; i < 4; i++) {
po0[i] = Converter::convert(
make_float2(elements[2 * i], elements[2 * i + 1]));
}
#pragma unroll
for (int i = 0; i < 4; i++) {
po1[i] = Converter::convert(
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
}
scalar_t_in* dst =
q_inout +
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
dim_base;
*reinterpret_cast<uint4*>(dst) = out0;
*reinterpret_cast<uint4*>(dst + 8) = out1;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
return;
}
// ═══════════════════════════════════════════════════════════════════════
// KV branch.
// ═══════════════════════════════════════════════════════════════════════
int64_t const slot_id = slot_mapping[tokenIdx];
if (slot_id < 0) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
return;
}
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_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;
// Round K to bf16 first, matching the unfused reference path where K is
// materialized as bf16 before K quantization. absmax, clamp, and FP8
// quant below all run on these bf16-rounded values.
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
elements[i] = Converter::convert(Converter::convert(elements[i]));
}
// Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
// shuffle requires full participation). RoPE lanes contribute garbage,
// but their values are gated out below via `!is_rope_lane`.
float local_absmax = 0.0f;
#pragma unroll
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 exponent = ceilf(log2f(absmax / kFp8Max));
float const inv_scale = exp2f(-exponent);
if (!is_rope_lane) {
// ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
uint8_t out_bytes[kElemsPerLane];
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
float scaled = elements[i] * inv_scale;
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
__nv_fp8_storage_t s =
__nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
out_bytes[i] = static_cast<uint8_t>(s);
}
// One 16-byte STG per lane.
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
*reinterpret_cast<uint4 const*>(out_bytes);
// Lane (4k) of each 4-lane group writes the scale byte for block k<7.
if ((laneId & 3) == 0) {
int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes
float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
}
// Lane 0 also writes the padding byte at index 7.
if (laneId == 0) {
token_scale_ptr[kNumQuantBlocks] = 0; // pad
}
} else {
// ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
uint4 out0, out1;
typename Converter::packed_hip_type* po0 =
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
typename Converter::packed_hip_type* po1 =
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
#pragma unroll
for (int i = 0; i < 4; i++) {
po0[i] = Converter::convert(
make_float2(elements[2 * i], elements[2 * i + 1]));
}
#pragma unroll
for (int i = 0; i < 4; i++) {
po1[i] = Converter::convert(
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
}
int const rope_local_base = dim_base - kNopeDim; // in [0, 64)
scalar_t_in* bf16_dst =
reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
*reinterpret_cast<uint4*>(bf16_dst) = out0;
*reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
}
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
}
#endif
}
// ────────────────────────────────────────────────────────────────────────────
// Launch wrapper
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t_in>
void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
scalar_t_in* q_inout, scalar_t_in const* kv_in, 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,
cudaStream_t stream) {
constexpr int kBlockSize = 256;
constexpr int kWarpsPerBlock = kBlockSize / 32;
int64_t const total_warps =
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
int const grid =
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
// PDL: enable programmatic stream serialization whenever the hardware
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
// so leave numAttrs = 0 and launch as a regular kernel.
static int const sm_version = getSMVersion();
// Host-side guard: the device kernel body is compiled as a no-op for
// bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
// unavailable there. Refuse the launch loudly instead of silently
// skipping the work.
TORCH_CHECK(
sm_version >= 80,
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
"(Ampere or newer); got sm_",
sm_version);
cudaLaunchConfig_t config;
config.gridDim = dim3(grid);
config.blockDim = dim3(kBlockSize);
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
cudaLaunchKernelEx(
&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
q_inout, 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);
}
} // namespace deepseek_v4_fused_ops
} // namespace vllm
// ────────────────────────────────────────────────────────────────────────────
// Torch op wrapper
// ────────────────────────────────────────────────────────────────────────────
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
torch::Tensor& q, // [N, H, 512] bf16, in place
torch::Tensor const& kv, // [N, 512] bf16 (read-only)
torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
torch::Tensor const& slot_mapping, // [N] int64
torch::Tensor const& position_ids, // [N] int64
torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
double eps, int64_t cache_block_size) {
TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
"slot_mapping must be int64 CUDA");
TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
"position_ids must be int64 CUDA");
TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
"cos_sin_cache shape [max_pos, 64]");
TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
"cos_sin_cache must be float32");
// With DP padding, slot_mapping can be shorter than q/kv/positions.
// Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
// KV quant+insert runs only on the first slot_mapping.size(0) rows.
int const num_tokens_full = static_cast<int>(q.size(0));
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
static_cast<int>(position_ids.size(0)) == num_tokens_full,
"q/kv/position_ids row counts must match");
TORCH_CHECK(num_tokens_insert <= num_tokens_full,
"slot_mapping must not exceed q row count");
int const num_heads_q = static_cast<int>(q.size(1));
int const cache_block_size_i = static_cast<int>(cache_block_size);
int const kv_block_stride = static_cast<int>(k_cache.stride(0));
at::cuda::OptionalCUDAGuard device_guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream();
VLLM_DISPATCH_HALF_TYPES(
q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
using qkv_scalar_t = scalar_t;
vllm::deepseek_v4_fused_ops::
launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
num_tokens_full, num_tokens_insert, num_heads_q,
cache_block_size_i, kv_block_stride, stream);
});
}
+7 -15
View File
@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
}
v_out[i] = dst;
}
@@ -135,17 +134,10 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
_f16Vec<scalar_t, width> res = residual_v[id];
_f16Vec<scalar_t, width> w = weight_v[idx];
_f16Vec<scalar_t, width> out;
using Converter = _typeConvert<scalar_t>;
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
}
input_v[strided_id] = out;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
input_v[strided_id] = temp;
}
}
@@ -182,8 +174,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
input[blockIdx.x * input_stride + idx] =
((scalar_t)(x * s_variance)) * weight[idx];
}
}
+10 -28
View File
@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
}
}
}
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
_f16Vec<scalar_t, width> res = residual_v[id];
_f16Vec<scalar_t, width> w = weight_v[idx];
using Converter = _typeConvert<scalar_t>;
using HipT = typename Converter::hip_type;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
out[id * width + i] =
scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
}
}
}
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
out[blockIdx.x * hidden_size + idx] =
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
}
}
@@ -277,9 +277,7 @@ void quant_impl(void* output, void* output_scale, void* input,
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
if (blockRepeat > 1) {
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
// The shared-memory vectorized offset load only handles full 4-expert
// chunks. Use the scalar specialization for the remainder cases.
if (n_experts >= 4 && n_experts % 4 == 0) {
if (n_experts >= 4) {
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
@@ -301,9 +299,7 @@ void quant_impl(void* output, void* output_scale, void* input,
n_experts);
}
} else {
// The low-latency vectorized expert lookup only handles full 16-expert
// chunks. Fall back to the scalar lookup path for the remainder cases.
if (n_experts >= 16 && n_experts % 16 == 0) {
if (n_experts >= 16) {
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
<<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
-9
View File
@@ -12,15 +12,6 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
torch::Tensor& gating_output, bool renormalize,
std::optional<torch::Tensor> bias);
void topk_softplus_sqrt(torch::Tensor& topk_weights,
torch::Tensor& topk_indices,
torch::Tensor& token_expert_indices,
torch::Tensor& gating_output, bool renormalize,
double routed_scaling_factor,
const c10::optional<torch::Tensor>& correction_bias,
const c10::optional<torch::Tensor>& input_ids,
const c10::optional<torch::Tensor>& tid2eid);
void moe_sum(torch::Tensor& input, torch::Tensor& output);
void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
-715
View File
@@ -1,715 +0,0 @@
/*
* Adapted from
* https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
* Copyright (c) 2024, The vLLM team.
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <type_traits>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "../cuda_compat.h"
#include "../cub_helpers.h"
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#else
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
typedef __hip_bfloat16 __nv_bfloat16;
typedef __hip_bfloat162 __nv_bfloat162;
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
namespace vllm {
namespace moe {
/// Aligned array type
template <typename T,
/// Number of elements in the array
int N,
/// Alignment requirement in bytes
int Alignment = sizeof(T) * N>
struct alignas(Alignment) AlignedArray {
T data[N];
};
template <typename T>
__device__ __forceinline__ float toFloat(T value) {
if constexpr (std::is_same_v<T, float>) {
return value;
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
return __bfloat162float(value);
} else if constexpr (std::is_same_v<T, __half>) {
return __half2float(value);
}
}
#define FINAL_MASK 0xffffffff
template <typename T>
__inline__ __device__ T warpReduceSum(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
return val;
}
// ====================== TopK softplus_sqrt things
// ===============================
/*
A Top-K gating softplus_sqrt written to exploit when the number of experts in
the MoE layers are a small power of 2. This allows us to cleanly share the
rows among the threads in a single warp and eliminate communication between
warps (so no need to use shared mem).
It fuses the sigmoid, max and argmax into a single kernel.
Limitations:
1) This implementation is optimized for when the number of experts is a small
power of 2. Additionally it also supports when number of experts is multiple
of 64 which is still faster than the computing sigmoid and topK separately
(only tested on CUDA yet). 2) This implementation assumes k is small, but will
work for any k.
*/
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
typename InputType = float>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGatingSoftplusSqrt(
const InputType* input, const bool* finished, float* output,
const int num_rows, IndType* indices, int* source_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const IndType* input_ids, const IndType* tid2eid) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
"InputType must be float, __nv_bfloat16, or __half");
// We begin by enforcing compile time assertions and setting up compile time
// constants.
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
"BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>) {
static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
"ELTS_PER_LDG must be 1 or even for 16-bit conversion");
}
// Restrictions based on previous section.
static_assert(
VPT % ELTS_PER_LDG == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
"THREADS_PER_ROW must be power of 2");
static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
"THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
// Restrictions for previous section.
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
"The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time
// variables. ========================
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a
// block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block.
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
// Now, using the base row per thread block, we compute the base row per warp.
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
// The threads in a warp are split into sub-groups that will work on a row.
// We compute row offset for each thread sub-group
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
if (thread_row >= num_rows) {
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
// Now, we compute the group each thread belong to in order to determine the
// first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Finally, we pull in the data from global mem
float row_chunk[VPT];
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
#endif
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr =
reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
}
} else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
if constexpr (ELTS_PER_LDG >= 2) {
using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
const VecType* vec_thread_read_ptr =
reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
#pragma unroll
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
*reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
}
}
} else { // ELTS_PER_LDG == 1
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
const __nv_bfloat16* scalar_ptr =
thread_read_ptr + ii * THREADS_PER_ROW;
row_chunk[ii] = __bfloat162float(*scalar_ptr);
}
}
} else if constexpr (std::is_same_v<InputType, __half>) {
if constexpr (ELTS_PER_LDG >= 2) {
using VecType = AlignedArray<__half, ELTS_PER_LDG>;
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
const VecType* vec_thread_read_ptr =
reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
#pragma unroll
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
row_chunk_f2[base_idx_f2 + jj] = __half22float2(
*reinterpret_cast<const __half2*>(vec.data + jj * 2));
}
}
} else { // ELTS_PER_LDG == 1
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
row_chunk[ii] = __half2float(*scalar_ptr);
}
}
}
constexpr float threshold = 20.0f;
constexpr float beta = 1.0f;
// Hash MoE path: indices are predetermined from lookup table
if constexpr (USE_HASH) {
const IndType token_id = input_ids[thread_row];
const IndType* expert_indices_for_token = tid2eid + token_id * k;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
row_chunk[ii] = sqrtf(val);
}
float selected_sum = 0.f;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int expert = expert_indices_for_token[k_idx];
const int idx = k * thread_row + k_idx;
for (int ii = 0; ii < VPT; ++ii) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
indices[idx] = expert;
selected_sum += row_chunk[ii];
break;
}
}
}
// Compute per-thread scale (using warp reduction when renormalizing).
if (renormalize) {
selected_sum = warpReduceSum(selected_sum);
}
float scale = static_cast<float>(routed_scaling_factor);
if (renormalize) {
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
scale /= denom;
}
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int expert = expert_indices_for_token[k_idx];
const int idx = k * thread_row + k_idx;
for (int ii = 0; ii < VPT; ++ii) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
output[idx] = row_chunk[ii] * scale;
break;
}
}
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;");
#endif
return;
}
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
// Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
// the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread;
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
float selected_sum = 0.f;
for (int k_idx = 0; k_idx < k; ++k_idx) {
// First, each thread does the local argmax
float max_val = row_chunk[0];
int expert = start_col;
#pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
++ldg, col += COLS_PER_GROUP_LDG) {
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
// No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=)
if (val > max_val) {
max_val = val;
expert = col + ii;
}
}
}
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
// reach consensus about the max. This will be useful for K > 1 so that the
// threads can agree on "who" had the max value. That thread can then blank out
// their max with -inf and the warp can run more iterations...
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
float other_max =
VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
int other_expert =
VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
// We want lower indices to "win" in every thread so we break ties this
// way
if (other_max > max_val ||
(other_max == max_val && other_expert < expert)) {
max_val = other_max;
expert = other_expert;
}
}
// Write the max for this k iteration to global memory.
if (thread_group_idx == 0) {
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results to
// global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
}
}
// Finally, we clear the value in the thread with the current max if there
// is another iteration to run.
if (k_idx + 1 < k) {
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
const int thread_to_clear_in_group =
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
// Only the thread in the group which produced the max will reset the
// "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) {
const int offset_for_expert = expert % ELTS_PER_LDG;
// Safe to set to any negative value since row_chunk values must be
// between 0 and 1.
row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
-10000.f;
}
}
}
// Apply renormalization and routed scaling factor to final weights.
if (thread_group_idx == 0) {
float scale = static_cast<float>(routed_scaling_factor);
if (renormalize) {
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
scale /= denom;
}
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] * scale;
}
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;");
#endif
}
namespace detail {
// Constructs some constants needed to partition the work across threads at
// compile time.
template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
typename InputType>
struct TopkConstants {
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
"");
static constexpr int VECs_PER_THREAD =
MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
};
} // namespace detail
#define DISPATCH_HASH(use_hash, USE_HASH, ...) \
if (use_hash) { \
const bool USE_HASH = true; \
static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
__VA_ARGS__ \
} else { \
const bool USE_HASH = false; \
static_assert(USE_HASH == false, \
"USE_HASH must be compile-time constant"); \
__VA_ARGS__ \
}
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
void topkGatingSoftplusSqrtLauncherHelper(
const InputType* input, const bool* finished, float* output,
IndType* indices, int* source_row, const int num_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
cudaStream_t stream) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
static constexpr int VPT = Constants::VPT;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
DISPATCH_HASH(use_hash, USE_HASH, {
auto* kernel =
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
#ifndef USE_ROCM
cudaLaunchConfig_t config = {};
config.gridDim = num_blocks;
config.blockDim = block_dim;
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.numAttrs = 1;
config.attrs = attrs;
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
indices, source_row, k, start_expert, end_expert,
renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
#else
kernel<<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert,
end_expert, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
#endif
})
}
#ifndef USE_ROCM
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
static_assert(WARP_SIZE == 32, \
"Unsupported warp size. Only 32 is supported for CUDA"); \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
stream);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
} else if (WARP_SIZE == 32) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
}
#endif
template <typename IndType, typename InputType>
void topkGatingSoftplusSqrtKernelLauncher(
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
int* token_expert_indices, const int num_tokens, const int num_experts,
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash, const IndType* input_ids,
const IndType* tid2eid, cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
// elements can be loaded by a warp
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
(std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>)
? 4
: 8;
#endif
switch (num_experts) {
case 1:
LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 2:
LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 4:
LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 8:
LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 16:
LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 32:
LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 64:
LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 128:
LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 256:
LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 512:
LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
// num_experts, alternatively we can test 4 bytes loading and enable it in
// future.
#ifndef USE_ROCM
case 192:
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 320:
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 384:
LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 448:
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 576:
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
#endif
default: {
TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
}
}
}
} // namespace moe
} // namespace vllm
template <typename ComputeType>
void dispatch_topk_softplus_sqrt_launch(
const ComputeType* gating_output, torch::Tensor& topk_weights,
torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
int num_tokens, int num_experts, int topk, bool renormalize,
double routed_scaling_factor,
const c10::optional<torch::Tensor>& correction_bias,
const c10::optional<torch::Tensor>& input_ids,
const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().data_ptr<float>();
}
bool use_hash = false;
if (tid2eid.has_value()) {
TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
use_hash = true;
}
if (topk_indices.scalar_type() == at::ScalarType::Int) {
const int* input_ids_ptr = nullptr;
const int* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().data_ptr<int>();
tid2eid_ptr = tid2eid.value().data_ptr<int>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
gating_output, topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
} else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
const uint32_t* input_ids_ptr = nullptr;
const uint32_t* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
gating_output, topk_weights.data_ptr<float>(),
topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
} else {
TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
const int64_t* input_ids_ptr = nullptr;
const int64_t* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().data_ptr<int64_t>();
tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
gating_output, topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
}
}
void topk_softplus_sqrt(
torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& token_expert_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize, double routed_scaling_factor,
const c10::optional<torch::Tensor>& correction_bias,
const c10::optional<torch::Tensor>& input_ids,
const c10::optional<torch::Tensor>& tid2eid) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (gating_output.scalar_type() == at::ScalarType::Float) {
dispatch_topk_softplus_sqrt_launch<float>(
gating_output.data_ptr<float>(), topk_weights, topk_indices,
token_expert_indices, num_tokens, num_experts, topk, renormalize,
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() == at::ScalarType::Half) {
dispatch_topk_softplus_sqrt_launch<__half>(
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
} else {
TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
}
}
-8
View File
@@ -16,14 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
"bias) -> ()");
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
#ifndef USE_ROCM
m.def(
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, float "
"routed_scaling_factor, Tensor? "
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
#endif
// Calculate the result of moe by summing up the partial results
// from all selected experts.
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
+1 -9
View File
@@ -100,11 +100,6 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
bool is_neox, torch::Tensor& position_ids,
int64_t forced_token_heads_per_warp);
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
void apply_repetition_penalties_(torch::Tensor& logits,
const torch::Tensor& prompt_mask,
const torch::Tensor& output_mask,
@@ -158,13 +153,10 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
std::optional<torch::Tensor> key, int64_t head_size,
torch::Tensor& cos_sin_cache, bool is_neox,
int64_t rope_dim_offset, bool inverse);
torch::Tensor& cos_sin_cache, bool is_neox);
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
+16 -17
View File
@@ -18,6 +18,7 @@ namespace persistent {
// Constants
// ============================================================================
constexpr int TopK = 2048;
constexpr int kThreadsPerBlock = 1024;
constexpr int RADIX = 256;
@@ -127,12 +128,11 @@ struct RadixRowState {
struct PersistentTopKParams {
const float* __restrict__ input; // [num_rows, stride]
int32_t* __restrict__ output; // [num_rows, top_k]
int32_t* __restrict__ output; // [num_rows, TopK]
int32_t* __restrict__ lengths; // [num_rows]
RadixRowState* row_states; // large path: per-group state
uint32_t num_rows;
uint32_t stride;
uint32_t top_k; // actual k value for output stride
uint32_t chunk_size; // large path: elements per CTA
uint32_t ctas_per_group; // 1=medium, >1=large
uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit)
@@ -154,7 +154,6 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
return key >> 5;
}
template <int TopK>
__device__ __noinline__ void histogram_2048_topk(
const float* __restrict__ logits, int32_t* __restrict__ output_indices,
int32_t seq_len) {
@@ -419,7 +418,6 @@ __device__ __noinline__ void histogram_2048_topk(
// by: DarkSharpness
// which at the same time is an optimized topk kernel copied from tilelang
// kernel
template <int TopK>
__device__ __noinline__ void histogram_256_topk(
const float* __restrict__ logits, int* __restrict__ output_indices,
int logits_offset, int seq_len) {
@@ -651,7 +649,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
// ============================================================================
template <int TopK, uint32_t VEC_SIZE>
template <uint32_t VEC_SIZE>
__device__ void radix_topk(const float* __restrict__ row_input,
int32_t* __restrict__ row_output, uint32_t seq_len,
uint32_t my_chunk_start, uint32_t chunk_size,
@@ -859,7 +857,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
// see filtered_topk.cuh)
// ============================================================================
template <int TopK = 2048, uint32_t VEC_SIZE = 1>
template <uint32_t VEC_SIZE = 1>
__global__ void __launch_bounds__(kThreadsPerBlock, 2)
persistent_topk_kernel(PersistentTopKParams params) {
const uint32_t tx = threadIdx.x;
@@ -917,7 +915,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
if (row_idx >= params.num_rows) break;
const uint32_t seq_len = params.lengths[row_idx];
int32_t* row_output = params.output + row_idx * params.top_k;
int32_t* row_output = params.output + row_idx * TopK;
const float* row_input = params.input + row_idx * params.stride;
if (seq_len <= RADIX_THRESHOLD) {
@@ -929,19 +927,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
}
} else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
histogram_2048_topk<TopK>(row_input, row_output, seq_len);
histogram_2048_topk(row_input, row_output, seq_len);
} else {
histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
histogram_256_topk(row_input, row_output, 0, seq_len);
}
}
continue;
}
const uint32_t my_chunk_start = cta_in_group * chunk_size;
radix_topk<TopK, VEC_SIZE>(
row_input, row_output, seq_len, my_chunk_start, chunk_size,
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
cta_in_group, ctas_per_group, barrier_phase, iter, tx);
radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
chunk_size, local_histogram, suffix_sum,
shared_scalars, shared_ordered, state, cta_in_group,
ctas_per_group, barrier_phase, iter, tx);
}
}
@@ -1013,6 +1011,7 @@ struct FilteredTopKTraits<float> {
}
};
constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
16 * 1024; // 16K indices per buffer
@@ -1026,7 +1025,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
* \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>
template <typename DType, typename IdType, int VEC_SIZE>
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
IdType* __restrict__ output,
@@ -1060,7 +1059,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
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];
alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
auto& s_histogram = s_histogram_buf[0];
@@ -1281,7 +1280,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
return static_cast<int>(g);
}
template <typename DType, typename IdType, uint32_t MAX_K = 2048>
template <typename DType, typename IdType>
cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
IdType* lengths, uint32_t num_rows,
uint32_t top_k_val, uint32_t max_len,
@@ -1298,7 +1297,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
#define DISPATCH_VEC_SIZE(VS) \
if (vec_size == VS) { \
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>; \
FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
+33 -39
View File
@@ -9,29 +9,28 @@ namespace vllm {
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
const bool inverse) {
scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
int x_index, y_index;
float cos_f, sin_f;
scalar_t cos, sin;
if (IS_NEOX) {
// GPT-NeoX style rotary embedding.
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos_f = VLLM_LDG(cos_ptr + x_index);
sin_f = VLLM_LDG(sin_ptr + x_index);
cos = VLLM_LDG(cos_ptr + x_index);
sin = VLLM_LDG(sin_ptr + x_index);
} else {
// GPT-J style rotary embedding.
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos_f = VLLM_LDG(cos_ptr + x_index / 2);
sin_f = VLLM_LDG(sin_ptr + x_index / 2);
cos = VLLM_LDG(cos_ptr + x_index / 2);
sin = VLLM_LDG(sin_ptr + x_index / 2);
}
if (inverse) {
sin_f = -sin_f;
}
const float x_f = static_cast<float>(arr[x_index]);
const float y_f = static_cast<float>(arr[y_index]);
arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
const scalar_t x = arr[x_index];
const scalar_t y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;
}
template <typename scalar_t, bool IS_NEOX>
@@ -43,23 +42,22 @@ inline __device__ void apply_rotary_embedding(
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const float* cache_ptr, const int head_size, const int num_heads,
const scalar_t* cache_ptr, const int head_size, const int num_heads,
const int num_kv_heads, const int rot_dim, const int token_idx,
const int64_t query_stride, const int64_t key_stride,
const int64_t head_stride, const int64_t rope_dim_offset,
const bool inverse) {
const int64_t head_stride) {
const int embed_dim = rot_dim / 2;
const float* cos_ptr = cache_ptr;
const float* sin_ptr = cache_ptr + embed_dim;
const scalar_t* cos_ptr = cache_ptr;
const scalar_t* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
token_idx * query_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
if (key != nullptr) {
@@ -67,10 +65,10 @@ inline __device__ void apply_rotary_embedding(
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
token_idx * key_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
}
}
@@ -86,18 +84,19 @@ __global__ void rotary_embedding_kernel(
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
// 2]
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
const int64_t head_stride, const int num_heads, const int num_kv_heads,
const int head_size, const int64_t rope_dim_offset, const bool inverse) {
const int head_size) {
// Each thread block is responsible for one token.
const int token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const float* cache_ptr = cos_sin_cache + pos * rot_dim;
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, IS_NEOX>(
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
inverse);
token_idx, query_stride, key_stride, head_stride);
}
} // namespace vllm
@@ -116,7 +115,7 @@ void rotary_embedding(
// [num_tokens, num_heads, head_size]
int64_t head_size,
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
bool is_neox, int64_t rope_dim_offset, bool inverse) {
bool is_neox) {
// num_tokens = batch_size * seq_len
int64_t num_tokens = positions.numel();
int positions_ndim = positions.dim();
@@ -155,8 +154,6 @@ void rotary_embedding(
int seq_dim_idx = positions_ndim - 1;
int64_t query_stride = query.stride(seq_dim_idx);
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
// Determine head stride: for [*, heads, head_size] use stride of last dim;
// for flat [*, heads*head_size], heads blocks are contiguous of size
// head_size
@@ -168,23 +165,20 @@ void rotary_embedding(
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
if (is_neox) {
vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
inverse);
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
head_stride, num_heads, num_kv_heads, head_size);
} else {
vllm::rotary_embedding_kernel<scalar_t, false>
<<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
inverse);
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
key_stride, head_stride, num_heads, num_kv_heads, head_size);
}
});
}
+11 -2
View File
@@ -40,6 +40,15 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
#define __HIP__FP8MFMA__
#endif
#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
defined(__gfx1150__) || defined(__gfx1151__))
#define __HIP__GFX11__
#endif
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
#define __HIP__GFX12__
#endif
#if defined(NDEBUG)
#undef NDEBUG
#include <assert.h>
@@ -1620,7 +1629,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
}
}
#elif defined(__GFX11__)
#elif defined(__HIP__GFX11__)
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
@@ -2379,7 +2388,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
}
#elif defined(__GFX12__)
#elif defined(__HIP__GFX12__)
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
+23 -18
View File
@@ -26,11 +26,16 @@
#define __HIP__GFX9__
#endif
// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
#if defined(__GFX11__) || defined(__GFX12__)
#if defined(__HIPCC__) && \
(defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
#define __HIP__GFX1X__
#endif
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
#define __HIP__GFX12__
#endif
#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
#define __HIP__MI3XX__
#endif
@@ -1840,7 +1845,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
return out_c;
}
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
int A_CHUNK, int UNRL, int N>
__global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1888,7 +1893,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
float sB = *s_B;
while (m < M) {
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
float sum[N][YTILE] = {};
#else
@@ -1926,7 +1931,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#pragma unroll
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
for (uint32_t n = 0; n < N; n++) {
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
for (int y = 0; y < YTILE; ++y) {
#pragma unroll
@@ -1950,7 +1955,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
// Final reduction
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
for (int n = 0; n < N; n++) {
for (int y = 0; y < YTILE; y++) {
@@ -1988,7 +1993,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#endif
const bool writeback_lane =
#ifdef __GFX12__
#ifdef __HIP__GFX12__
threadIdx.x == (THRDS - 1);
#else
threadIdx.x == 0;
@@ -2004,7 +2009,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
for (int n = 0; n < N; n++) {
for (int y = 0; y < YTILE; y++) {
if (y + m >= M) break; // To avoid mem access fault.
#ifdef __GFX12__
#ifdef __HIP__GFX12__
float result = sum[n][y] * sA * sB;
#else
float result = sum[n][y][0] * sA * sB;
@@ -2022,7 +2027,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
m += CuCount * _WvPrGrp * YTILE;
}
}
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
int A_CHUNK, int UNRL, int N>
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
@@ -2034,9 +2039,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
const int _WvPrGrp, const int CuCount) {
UNREACHABLE_CODE
}
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
int A_CHUNK, int UNRL, int N>
__global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -2083,7 +2088,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
float sB = *s_B;
while (m < M) {
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
float sum[N][YTILE] = {};
#else
@@ -2123,7 +2128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#pragma unroll
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
for (uint32_t n = 0; n < N; n++) {
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
for (int y = 0; y < YTILE; ++y) {
#pragma unroll
@@ -2147,7 +2152,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
// Final reduction
#ifdef __GFX12__
#ifdef __HIP__GFX12__
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
for (int n = 0; n < N; n++) {
for (int y = 0; y < YTILE; y++) {
@@ -2185,7 +2190,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
#endif
const bool writeback_lane =
#ifdef __GFX12__
#ifdef __HIP__GFX12__
threadIdx.x == (THRDS - 1);
#else
threadIdx.x == 0;
@@ -2201,7 +2206,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
for (int n = 0; n < N; n++) {
for (int y = 0; y < YTILE; y++) {
if (y + m >= M) break; // To avoid mem access fault.
#ifdef __GFX12__
#ifdef __HIP__GFX12__
float result = sum[n][y] * sA * sB;
#else
float result = sum[n][y][0] * sA * sB;
@@ -2219,7 +2224,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
m += CuCount * _WvPrGrp * YTILE;
}
}
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
int A_CHUNK, int UNRL, int N>
__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
@@ -2231,7 +2236,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
const int CuCount) {
UNREACHABLE_CODE
}
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
+1 -7
View File
@@ -258,13 +258,7 @@ __device__ bool processHistogramStep(
auto processBins = [&](float logit, int idx) {
if (isPartialMatch<patternShift>(logit, logitPattern)) {
uint32_t binIdx = extractBinIdx<step>(logit);
// Only write elements with binIdx < thresholdBinIdx when:
// 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);
if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
if (binIdx < thresholdBinIdx) {
// The element is part of the top-k selection
int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
+38 -117
View File
@@ -10,17 +10,33 @@
#include "persistent_topk.cuh"
#endif
namespace {
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
int64_t max_seq_len) {
#ifndef USE_ROCM
template <int TopK>
void launch_persistent_topk(const torch::Tensor& logits,
const torch::Tensor& lengths, torch::Tensor& output,
torch::Tensor& workspace, int64_t max_seq_len) {
namespace P = vllm::persistent;
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
"lengths must be 1D or 2D");
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
TORCH_CHECK(output.dim() == 2, "output must be 2D");
const int64_t num_rows = logits.size(0);
const int64_t stride = logits.size(1);
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
"output size mismatch");
namespace P = vllm::persistent;
TORCH_CHECK(k == P::TopK, "k must be 2048");
TORCH_CHECK(k <= stride, "k out of range");
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
static int num_sms = 0;
@@ -34,17 +50,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
}
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
TORCH_CHECK(status == cudaSuccess,
"FilteredTopK failed: ", cudaGetErrorString(status));
} else {
TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
// Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
// large path. Empirically tuned.
int effective_max_smem;
if (num_rows <= 4) {
effective_max_smem =
@@ -82,73 +99,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
// Query occupancy for the instantiation that will actually launch;
// overestimating it deadlocks the cooperative barrier.
int occupancy = 1;
cudaError_t occ_err = cudaSuccess;
if (vec_size == 4) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
smem_size);
} else if (vec_size == 2) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
smem_size);
} else {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
smem_size);
}
TORCH_CHECK(occ_err == cudaSuccess,
"persistent_topk occupancy query failed: ",
cudaGetErrorString(occ_err));
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
smem_size);
if (occupancy < 1) occupancy = 1;
// The cooperative spin-wait barrier only runs when at least one row hits
// the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
// early-exit, so oversubscription can't deadlock and headroom is wasted.
const bool needs_cooperative =
static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
const uint32_t hw_resident_cap =
static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
uint32_t max_resident_ctas = hw_resident_cap;
if (needs_cooperative) {
// Reserve one CTA per SM when occupancy allows; fall back to a single
// CTA when occupancy == 1 (the most deadlock-prone case — any straggler
// kernel that takes the only slot on one SM hangs the barrier). Never
// drop below one full group's worth.
uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
if (max_resident_ctas >= headroom + ctas_per_group) {
max_resident_ctas -= headroom;
}
}
uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
static_cast<uint32_t>(num_rows));
if (num_groups == 0) num_groups = 1;
uint32_t total_ctas = num_groups * ctas_per_group;
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
// instead of deadlocking. Only relevant when needs_cooperative.
if (needs_cooperative && total_ctas > hw_resident_cap) {
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
"persistent_topk would oversubscribe and the FilteredTopK "
"fallback requires >=128KB smem per block (have ",
max_smem_per_block, "). total_ctas=", total_ctas,
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
", smem=", smem_size, ").");
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
stream);
TORCH_CHECK(status == cudaSuccess,
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
return;
}
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
"workspace too small, need ", state_bytes, " bytes");
@@ -159,16 +121,15 @@ void launch_persistent_topk(const torch::Tensor& logits,
params.lengths = lengths.data_ptr<int32_t>();
params.num_rows = static_cast<uint32_t>(num_rows);
params.stride = static_cast<uint32_t>(stride);
params.top_k = static_cast<uint32_t>(TopK);
params.chunk_size = chunk_size;
params.row_states =
reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
params.ctas_per_group = ctas_per_group;
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
#define LAUNCH_PERSISTENT(TOPK_VAL, VS) \
#define LAUNCH_PERSISTENT(VS) \
do { \
auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
auto kernel = &P::persistent_topk_kernel<VS>; \
cudaError_t err = cudaFuncSetAttribute( \
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
TORCH_CHECK(err == cudaSuccess, \
@@ -177,11 +138,11 @@ void launch_persistent_topk(const torch::Tensor& logits,
} while (0)
if (vec_size == 4) {
LAUNCH_PERSISTENT(TopK, 4);
LAUNCH_PERSISTENT(4);
} else if (vec_size == 2) {
LAUNCH_PERSISTENT(TopK, 2);
LAUNCH_PERSISTENT(2);
} else {
LAUNCH_PERSISTENT(TopK, 1);
LAUNCH_PERSISTENT(1);
}
#undef LAUNCH_PERSISTENT
}
@@ -189,46 +150,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
cudaError_t err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess,
"persistent_topk failed: ", cudaGetErrorString(err));
}
#endif
} // anonymous namespace
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
int64_t max_seq_len) {
#ifndef USE_ROCM
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
"lengths must be 1D or 2D");
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
TORCH_CHECK(output.dim() == 2, "output must be 2D");
const int64_t num_rows = logits.size(0);
const int64_t stride = logits.size(1);
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
"output size mismatch");
TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
"persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
if (k == 512) {
launch_persistent_topk<512>(logits, lengths, output, workspace,
max_seq_len);
} else if (k == 1024) {
launch_persistent_topk<1024>(logits, lengths, output, workspace,
max_seq_len);
} else {
launch_persistent_topk<2048>(logits, lengths, output, workspace,
max_seq_len);
}
#else
TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
#endif
+1 -21
View File
@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
// SwiGLU activation with input clamping.
ops.def(
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
"-> ()");
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -183,19 +177,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"int forced_token_heads_per_warp=-1) -> ()");
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
#ifndef USE_ROCM
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
// kernel launch.
ops.def(
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
"Tensor! q, Tensor kv, Tensor! k_cache, "
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
"float eps, int cache_block_size) -> ()");
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
#endif
// Apply repetition penalties to logits in-place
ops.def(
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -259,8 +240,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
" Tensor cos_sin_cache, bool is_neox, int "
"rope_dim_offset=0, bool inverse=False) -> ()");
" Tensor cos_sin_cache, bool is_neox) -> ()");
ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
// Quantization ops
+60 -27
View File
@@ -22,7 +22,7 @@
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
# =============================================================================
ARG CUDA_VERSION=13.0.2
ARG CUDA_VERSION=13.0.0
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION=22.04
@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# Explicitly set the list to avoid issues with torch 2.2
# See https://github.com/pytorch/pytorch/pull/123243
# From versions.json: .torch.cuda_arch_list
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
#################### BUILD BASE IMAGE ####################
@@ -478,6 +478,9 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL
ENV DEBIAN_FRONTEND=noninteractive
WORKDIR /vllm-workspace
@@ -487,35 +490,43 @@ WORKDIR /vllm-workspace
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
# Install Python (via uv / python-build-standalone) and system dependencies.
# This replaces the deadsnakes PPA, removing the build-time dependency on
# Launchpad and matching how the build-stage (`base`) installs Python.
# python-build-standalone bundles dev headers, the venv module, and
# python3-config, so the python3.X-dev / python3.X-venv apt packages
# are not needed.
# Install Python and system dependencies
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
software-properties-common \
curl \
sudo \
ffmpeg \
libsm6 \
libxext6 \
libgl1 \
&& if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
mkdir -p -m 0755 /etc/apt/keyrings ; \
curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
fi ; \
else \
for i in 1 2 3; do \
add-apt-repository -y ppa:deadsnakes/ppa && break || \
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
done ; \
fi \
&& apt-get update -y \
&& apt-get install -y --no-install-recommends \
python${PYTHON_VERSION} \
python${PYTHON_VERSION}-dev \
python${PYTHON_VERSION}-venv \
libibverbs-dev \
&& rm -rf /var/lib/apt/lists/* \
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
&& $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
&& rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
&& ln -s /opt/venv/bin/python3 /usr/bin/python3 \
&& ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
&& ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
&& ln -s /opt/venv/bin/pip /usr/bin/pip \
&& update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
&& rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
&& curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
&& python3 --version && python3 -m pip --version
# Activate virtual environment and add uv to PATH
ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
ENV VIRTUAL_ENV="/opt/venv"
# Install CUDA development tools for runtime JIT compilation
# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -529,9 +540,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
libcurand-dev-${CUDA_VERSION_DASH} \
libcublas-${CUDA_VERSION_DASH} \
# Required by fastsafetensors (fixes #20384)
libnuma-dev \
# numactl CLI for NUMA binding at runtime
numactl && \
libnuma-dev && \
# Fixes nccl_allocator requiring nccl.h at runtime
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -540,6 +549,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
rm -rf /var/lib/apt/lists/*
# Install uv for faster pip installs
RUN python3 -m pip install uv
# Environment for uv
ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -568,12 +580,33 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.8.post1
# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
ARG FLASHINFER_VERSION=0.6.7
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
&& flashinfer show-config \
&& flashinfer download-cubin
&& flashinfer show-config
# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
# which fails without internet. This step caches them at build time.
RUN python3 <<'PYEOF'
from flashinfer.jit import env as jit_env
from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
from flashinfer.artifacts import ArtifactPath, CheckSumHash
download_trtllm_headers(
'bmm',
jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
ArtifactPath.TRTLLM_GEN_BMM,
get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
)
print('FlashInfer TRTLLM BMM headers downloaded successfully')
PYEOF
# ============================================================
# OPENAI API SERVER DEPENDENCIES
@@ -729,7 +762,7 @@ ENV HF_XET_HIGH_PERFORMANCE 1
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
# Copy in the v1 package for testing (it isn't distributed yet)
COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1
COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
# Source code is used in the `python_only_compile.sh` test
# We hide it inside `src/` so that this source code
@@ -753,7 +786,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
ENV UV_HTTP_TIMEOUT=500
# install kv_connectors if requested
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
+6 -3
View File
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system numba==0.65.0
uv pip install --system numba==0.61.2
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/common.txt
@@ -217,13 +217,16 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
# build flashinfer for torch nightly from source around 10 mins
# release version: v0.6.8.post1
# release version: v0.6.7
# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
# todo(elainewy): cache flashinfer build result for faster build
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/uv \
echo "git clone flashinfer..." \
&& git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
&& git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
&& cd flashinfer \
&& git submodule update --init --recursive \
&& echo "finish git clone flashinfer..." \
+8 -67
View File
@@ -2,11 +2,6 @@
ARG REMOTE_VLLM="0"
ARG COMMON_WORKDIR=/app
ARG BASE_IMAGE=rocm/vllm-dev:base
# AMD NIC backend
ARG NIC_BACKEND=none
# AMD AINIC apt repo settings
ARG AINIC_VERSION=1.117.5
ARG UBUNTU_CODENAME=jammy
# Sccache configuration (only used in release pipeline)
ARG USE_SCCACHE
@@ -124,10 +119,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
# RIXL/UCX build stages
FROM base AS build_rixl
ARG RIXL_BRANCH="bf4a7214"
ARG RIXL_BRANCH="f33a5599"
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
ARG UCX_BRANCH="7009d7a1"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ARG UCX_BRANCH="da3fac2a"
ARG UCX_REPO="https://github.com/ROCm/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +160,7 @@ RUN cd /usr/local/src && \
--disable-doxygen-doc \
--enable-optimizations \
--enable-devel-headers \
--with-rocm=${ROCM_PATH} \
--with-rocm=/opt/rocm \
--with-verbs \
--with-dm \
--enable-mt && \
@@ -186,12 +181,7 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
ninja install
# Generate RIXL wheel
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
RUN cd /opt/rixl && \
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
contrib/build-wheel.sh && \
mkdir -p /app/install && \
RUN cd /opt/rixl && mkdir -p /app/install && \
./contrib/build-wheel.sh \
--output-dir /app/install \
--rocm-dir ${ROCM_PATH} \
@@ -226,47 +216,6 @@ RUN git clone ${DEEPEP_REPO} \
&& git checkout ${DEEPEP_BRANCH} \
&& python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
# not force users to rebuild the long-lived Dockerfile.rocm_base image.
FROM base AS mori_base
ARG NIC_BACKEND
ARG AINIC_VERSION
ARG UBUNTU_CODENAME
RUN /bin/bash -lc 'set -euo pipefail; \
echo "[MORI] Install MoRI proxy deps"; \
pip install --quiet --ignore-installed blinker && \
pip install --quiet quart msgpack aiohttp pyzmq; \
echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
\
# NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
# Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
case "${NIC_BACKEND}" in \
# default: mlx5
none) \
;; \
# AMD NIC
ainic) \
apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
> /etc/apt/sources.list.d/amdainic.list; \
apt-get update && apt-get install -y --no-install-recommends \
libionic-dev \
ionic-common \
; \
rm -rf /var/lib/apt/lists/*; \
;; \
# TODO: Add Broadcom bnxt packages/repos here later.
# bnxt) \
# echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
# ;; \
*) \
echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
exit 2; \
;; \
esac;'
# -----------------------
# vLLM wheel release build stage (for building distributable wheels)
# This stage pins dependencies to custom ROCm wheel versions and handles version detection
@@ -369,7 +318,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
# -----------------------
# Test vLLM image
FROM mori_base AS test
FROM base AS test
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
@@ -436,10 +385,6 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
ENV MIOPEN_DEBUG_CONV_DIRECT=0
ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
# Source code is used in the `python_only_compile.sh` test
# We hide it inside `src/` so that this source code
# will not be imported by other tests
@@ -464,7 +409,7 @@ RUN printf '%s\n' \
# -----------------------
# Final vLLM image
FROM mori_base AS final
FROM base AS final
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
@@ -501,8 +446,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
ARG COMMON_WORKDIR
ARG BASE_IMAGE
ARG NIC_BACKEND
ARG AINIC_VERSION
# Copy over the benchmark scripts as well
COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
@@ -520,9 +463,7 @@ ENV HIP_FORCE_DEV_KERNARG=1
# Workaround for ROCm profiler limits
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
&& echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
&& echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt
CMD ["/bin/bash"]
+1 -1
View File
@@ -11,7 +11,7 @@ ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.10.post3"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_BRANCH="2d02c6a9"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
# Sccache configuration (only used in release pipeline)
+4 -4
View File
@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
rm "${ONECCL_INSTALLER}" && \
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# FIX triton
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.7.0
uv pip install triton-xpu==3.6.0
# remove torch bundled oneccl to avoid conflicts
RUN --mount=type=cache,target=/root/.cache/uv \
+3 -1
View File
@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
}
variable "TORCH_CUDA_ARCH_LIST" {
default = "8.0 8.9 9.0 10.0 11.0 12.0"
default = "8.0 8.9 9.0 10.0"
}
variable "COMMIT" {
@@ -88,6 +88,7 @@ target "test-ubuntu2404" {
args = {
UBUNTU_VERSION = "24.04"
GDRCOPY_OS_VERSION = "Ubuntu24_04"
FLASHINFER_AOT_COMPILE = "true"
}
output = ["type=docker"]
}
@@ -99,6 +100,7 @@ target "openai-ubuntu2404" {
args = {
UBUNTU_VERSION = "24.04"
GDRCOPY_OS_VERSION = "Ubuntu24_04"
FLASHINFER_AOT_COMPILE = "true"
}
output = ["type=docker"]
}
+5 -5
View File
@@ -2,7 +2,7 @@
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
"variable": {
"CUDA_VERSION": {
"default": "13.0.2"
"default": "13.0.0"
},
"PYTHON_VERSION": {
"default": "3.12"
@@ -11,10 +11,10 @@
"default": "22.04"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
"default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
},
"FINAL_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
"default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
},
"GET_PIP_URL": {
"default": "https://bootstrap.pypa.io/get-pip.py"
@@ -32,7 +32,7 @@
"default": "false"
},
"TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
"default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
},
"MAX_JOBS": {
"default": "2"
@@ -65,7 +65,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.8.post1"
"default": "0.6.7"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
Binary file not shown.

Before

Width:  |  Height:  |  Size: 315 KiB

After

Width:  |  Height:  |  Size: 315 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 108 KiB

File diff suppressed because one or more lines are too long
-32
View File
@@ -108,38 +108,6 @@ P99 ITL (ms): 8.39
==================================================
```
#### Results Visualization
The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
```bash
vllm bench serve \
--backend vllm \
--model meta-llama/Llama-3.1-8B-Instruct \
--endpoint /v1/completions \
--dataset-name sharegpt \
--dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
--num-prompts 100 \
--plot-timeline \
--timeline-itl-thresholds 2,5 \
--plot-dataset-stats \
--save-result
```
##### Interactive Timeline
The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
Example output:
<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
##### Dataset statistics
The generated figure shows the input prompt and output tokens distribution.
Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_dataset_stats.png)
#### Custom Dataset
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
+3 -3
View File
@@ -155,9 +155,9 @@ switch to `--physcpubind=<cpu-list> --membind=<node>`.
These `--numa-bind*` options only apply to GPU execution processes. They do not
configure the CPU backend's separate thread-affinity controls. Automatic
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as
ROCM-based platforms; other GPU backends must provide explicit binding lists if
they use these options.
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based platforms;
other GPU backends must provide explicit binding lists if they use these
options.
`--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in
the same order as the GPU indices.
+20 -19
View File
@@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model:
See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
#### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)
@@ -75,20 +75,21 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt
??? code "get_generation_prompt()"
```python
from vllm.config.speech_to_text import SpeechToTextParams
class YourASRModel(nn.Module, SupportsTranscription):
...
@classmethod
def get_generation_prompt(
cls,
stt_params: SpeechToTextParams,
audio: np.ndarray,
stt_config: SpeechToTextConfig,
model_config: ModelConfig,
language: str | None,
task_type: Literal["transcribe", "translate"],
request_prompt: str,
to_language: str | None,
) -> PromptType:
audio = stt_params.audio
stt_config = stt_params.stt_config
task_type = stt_params.task_type
# Example with a free-form instruction prompt
task_word = "Transcribe" if task_type == "transcribe" else "Translate"
prompt = (
"<start_of_turn>user\n"
@@ -111,22 +112,20 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:
??? code "get_generation_prompt()"
```python
from vllm.config.speech_to_text import SpeechToTextParams
class YourASRModel(nn.Module, SupportsTranscription):
...
@classmethod
def get_generation_prompt(
cls,
stt_params: SpeechToTextParams,
audio: np.ndarray,
stt_config: SpeechToTextConfig,
model_config: ModelConfig,
language: str | None,
task_type: Literal["transcribe", "translate"],
request_prompt: str,
to_language: str | None,
) -> PromptType:
audio = stt_params.audio
stt_config = stt_params.stt_config
language = stt_params.language
task_type = stt_params.task_type
request_prompt = stt_params.request_prompt
if language is None:
raise ValueError("Language must be specified")
@@ -214,13 +213,15 @@ Relevant server logic:
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
prompts = []
for chunk in chunks:
stt_params = request.build_stt_params(
prompt = self.model_cls.get_generation_prompt(
audio=chunk,
stt_config=self.asr_config,
model_config=self.model_config,
language=language,
task_type=self.task_type,
request_prompt=request.prompt,
to_language=to_language,
)
prompt = self.model_cls.get_generation_prompt(stt_params)
prompts.append(prompt)
return prompts, duration
```
-44
View File
@@ -4,7 +4,6 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
- [Deployment with CPUs](#deployment-with-cpus)
- [Deployment with GPUs](#deployment-with-gpus)
- [Serving with gRPC](#serving-with-grpc)
- [Troubleshooting](#troubleshooting)
- [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated)
- [Conclusion](#conclusion)
@@ -388,49 +387,6 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
If the service is correctly deployed, you should receive a response from the vLLM model.
## Serving with gRPC
vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies:
```bash
pip install vllm[grpc]
```
When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24).
To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes:
```yaml
containers:
- name: mistral-7b
image: vllm/vllm-openai:latest
command: ["/bin/sh", "-c"]
args: [
"pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code"
]
ports:
- containerPort: 50051
livenessProbe:
grpc:
port: 50051
initialDelaySeconds: 120
periodSeconds: 10
readinessProbe:
grpc:
port: 50051
initialDelaySeconds: 120
periodSeconds: 5
```
!!! note
The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`.
You can also verify the health service manually with `grpcurl`:
```bash
grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
```
## Troubleshooting
### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
+4 -4
View File
@@ -169,10 +169,10 @@ Priority is **1 = highest** (tried first).
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -213,9 +213,9 @@ configuration.
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
+6 -4
View File
@@ -52,14 +52,14 @@ For each graph replay:
When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
### Video inference support
### Video inference support (experimental)
Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
!!! note
Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
Mixed inputs (image+video) per prompt are also supported now.
Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
## Model integration via `SupportsEncoderCudaGraph`
@@ -76,7 +76,6 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
* `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
* `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
* `get_input_modality(...)` - return the modality of the inputs.
* `get_max_frames_per_video()` - return model-specific max frames per video.
!!! note
The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -97,7 +96,7 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
* `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
* `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
## Usage guide
@@ -143,6 +142,7 @@ Enable encoder CUDA Graphs via `compilation_config`:
```bash
vllm serve Qwen/Qwen3-VL-32B \
--limit-mm-per-prompt '{"image": 0}' \
--compilation-config '{"cudagraph_mm_encoder": true}'
```
@@ -150,6 +150,7 @@ With explicit budgets:
```bash
vllm serve Qwen/Qwen3-VL-32B \
--limit-mm-per-prompt '{"image": 0}' \
--compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
```
@@ -168,6 +169,7 @@ compilation_config = {
model = vllm.LLM(
model="Qwen/Qwen3-VL-32B",
limit_mm_per_prompt='{"image": 0}',
compilation_config=compilation_config,
)
```
+2 -42
View File
@@ -31,7 +31,6 @@ or just on the low or high end.
| [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
| [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
| [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always |
| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm → single kernel | O1 (ROCm/AITER only) | ~2% | No | Always |
## Support Matrix
@@ -44,7 +43,7 @@ The table below lists the quantization schemes supported by each fusion on each
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
@@ -52,7 +51,6 @@ The table below lists the quantization schemes supported by each fusion on each
| `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group |
| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group |
| `fuse_act_padding` | — | — | — | — | FP16/BF16 |
| `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 |
\* `fuse_attn_quant` support depends on the attention backend in use; not all backends support
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
@@ -152,7 +150,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt
- `FLASHINFER`: CUDA sm100+ with FlashInfer installed
`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
@@ -383,44 +381,6 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled.
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`)
### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`)
!!! info
ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention.
!!! note
When the native implementation of `rms_norm` is used (the default on CUDA and
ROCm for now), Inductor's built-in fusion already handles merging these norms
automatically. This explicit pass targets the case where AITER's custom
`rms_norm` op is active, which Inductor cannot fuse on its own.
**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm
operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER,
reducing kernel launch overhead from 2 launches to 1 per MLA layer.
```text
# Unfused:
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
q_c = rms_norm(q_c, q_weight, eps)
kv_c = rms_norm(kv_c, kv_weight, eps)
# Fused:
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
q_normed, kv_normed = fused_mla_dual_rms_norm(
q_c, q_weight, kv_c, kv_weight, eps1, eps2)
```
Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
when AITER is available.
**Code locations.**
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)
- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)
- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)
## See Also
- [Optimization Levels](optimization_levels.md) — high-level presets that set
+3 -3
View File
@@ -36,7 +36,7 @@ th {
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
-1
View File
@@ -56,7 +56,6 @@ Fusions:
- `-cc.pass_config.fuse_norm_quant=True`*
- `-cc.pass_config.fuse_act_quant=True`*
- `-cc.pass_config.fuse_act_padding=True`
- `-cc.pass_config.fuse_mla_dual_rms_norm=True`
\* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.</br>
† These fusions are ROCm-only and require AITER.
+1 -1
View File
@@ -104,7 +104,7 @@ for output in outputs:
Batch invariance has been tested and verified on the following models:
- **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
-70
View File
@@ -1,70 +0,0 @@
# Context Extension
!!! note
The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
This directory contains examples for extending the context length of models using vLLM.
## Offline Inference Example
The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
### Usage
```bash
python examples/offline_inference/context_extension.py
```
## OpenAI Online Method
You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
### Usage
Run the vLLM server with the following command to extend the context length using YARN:
```bash
vllm serve Qwen/Qwen3-0.6B \
--hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
--max-model-len 131072
```
### Client Example
After starting the server, you can use the OpenAI Python client to interact with it:
```python
from openai import OpenAI
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="token-abc123" # Dummy API key, required by the client
)
response = client.chat.completions.create(
model="Qwen/Qwen3-0.6B",
messages=[
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": "Hello"}
],
max_tokens=128,
temperature=0.8,
top_p=0.95
)
print(response.choices[0].message.content)
```
### Key Parameters
The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
Common parameters include:
- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
- `factor`: The factor by which to extend the context length
- `original_max_position_embeddings`: The original maximum position embeddings of the model
The following parameters are specific to vLLM:
- `max_model_len`: The new maximum sequence length after extension (original * factor).
Used for KV cache preallocation and request limit at serving time.
-64
View File
@@ -780,70 +780,6 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
Works with common video formats like MP4 when using OpenCV backends.
#### Pre-extracted Frame Sequences with `media_io_kwargs`
When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction.
**Supported Parameters:**
| Parameter | Type | Description |
| --------- | ---- | ----------- |
| `fps` | float | Frame rate of the original video |
| `frames_indices` | list[int] | Indices of the actually sampled frames |
| `total_num_frames` | int | Total frame count of the original video |
| `duration` | float | Duration of the original video in seconds |
| `do_sample_frames` | bool | Whether to perform frame sampling |
??? code
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
# Client-side frame extraction
frames = extract_frames(video_path, num_frames=32)
frames_b64 = ",".join([encode_image(f) for f in frames])
video_url = f"data:video/jpeg;base64,{frames_b64}"
# Pass video metadata via media_io_kwargs
response = client.chat.completions.create(
model="your-multimodal-model",
messages=[{
"role": "user",
"content": [
{"type": "video_url", "video_url": {"url": video_url}},
{"type": "text", "text": "Describe what happens in this video."}
]
}],
extra_body={
"media_io_kwargs": {
"video": {
"fps": 30.0,
"frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90,
100, 110, 120, 130, 140, 150, 160, 170,
180, 190, 200, 210, 220, 230, 240, 250,
260, 270, 280, 290, 300, 310],
"total_num_frames": 900,
"duration": 30.0,
}
}
},
)
print(response.choices[0].message.content)
```
**Why use `media_io_kwargs`?**
When extracting frames client-side, the server loses important context about the original video:
- **Temporal information**: Which frames were sampled and their positions in the original timeline
- **Video duration**: How long the original video was
- **Frame rate**: The original playback speed
By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped.
#### Custom RGBA Background Color
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
+2 -5
View File
@@ -17,9 +17,6 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
# MXFP8 quantization for weights and activations
llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8")
```
Or with the CLI:
@@ -27,7 +24,6 @@ Or with the CLI:
```bash
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
```
## Supported Schemes
@@ -36,7 +32,8 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
| ------ | ------------- | ------------------ | ----- |
| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback |
Support for additional schemes will be added in future versions of vllm.
## Advanced Configuration
+3 -1
View File
@@ -283,7 +283,9 @@ curl http://localhost:8000/v1/chat/completions \
"messages": [
{ "role": "user", "content": "9.11 and 9.8, which is greater?" }
],
"thinking_token_budget": 10
"extra_body": {
"thinking_token_budget": 10
}
}'
```
@@ -35,99 +35,6 @@ For reproducible measurements in your environment, use
[`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
or the [benchmark CLI guide](../../benchmarking/cli.md).
## `--speculative-config` schema
Use `--speculative-config` to pass speculative decoding settings as a JSON
object on the CLI:
```bash
vllm serve <target-model> \
--speculative-config '{
"method": "draft_model",
"model": "<draft-model>",
"num_speculative_tokens": 5
}'
```
The same keys are accepted from Python via `LLM(..., speculative_config={...})`.
The tables below highlight common user-facing keys accepted in this JSON
object; they are not an exhaustive schema reference.
For more details, see the generated [engine arguments reference](../../configuration/engine_args.md)
and the API docs for [vllm.config.SpeculativeConfig][].
### Common keys
These keys are commonly used across speculative decoding setups, though some
only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
`dflash`.
| Key | Type | Default | Allowed values / meaning |
| --- | --- | --- | --- |
| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. |
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
### Method-specific keys
#### N-gram
| Key | Type | Default | Meaning |
| --- | --- | --- | --- |
| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. |
| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. |
Example:
```bash
vllm serve <target-model> \
--speculative-config '{
"method": "ngram",
"num_speculative_tokens": 4,
"prompt_lookup_min": 2,
"prompt_lookup_max": 5
}'
```
#### Suffix decoding
| Key | Type | Default | Meaning |
| --- | --- | --- | --- |
| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. |
| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. |
| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. |
| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. |
Example:
```bash
vllm serve <target-model> \
--speculative-config '{
"method": "suffix",
"num_speculative_tokens": 8,
"suffix_decoding_max_tree_depth": 24,
"suffix_decoding_max_cached_requests": 10000,
"suffix_decoding_max_spec_factor": 1.0,
"suffix_decoding_min_token_prob": 0.1
}'
```
### Notes
- `--speculative-config` expects a JSON object on the CLI. In YAML config
files, use a nested mapping instead of an escaped JSON string.
- `tensor_parallel_size` is not a valid key in `speculative_config`. Use
`draft_tensor_parallel_size` instead.
- Keys such as `temperature` and `top_p` are sampling parameters, not
`--speculative-config` fields.
- Internal fields such as `target_model_config`, `draft_model_config`,
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
are populated by vLLM and are not intended to be set by users.
## Lossless guarantees of Speculative Decoding
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
@@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \
--port 8000 \
--seed 42 \
-tp 1 \
--max-model-len 2048 \
--gpu-memory-utilization 0.8 \
--speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
--max_model_len 2048 \
--gpu_memory_utilization 0.8 \
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
```
The code used to request as completions as a client remains unchanged:
@@ -77,8 +77,4 @@ The code used to request as completions as a client remains unchanged:
```
!!! warning
Note: Please use `--speculative-config` to set all configurations related
to speculative decoding. The previous method of specifying the model
through `--speculative-model` and adding related parameters such as
`--num-speculative-tokens` separately has been deprecated. For supported
keys and examples, see the [`--speculative-config` schema](README.md#--speculative-config-schema).
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
+1 -1
View File
@@ -38,7 +38,7 @@ for output in outputs:
```bash
vllm serve XiaomiMiMo/MiMo-7B-Base \
--tensor-parallel-size 1 \
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
--speculative_config '{"method":"mtp","num_speculative_tokens":1}'
```
## Notes
@@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \
--port 8000 \
--seed 42 \
-tp 1 \
--max-model-len 2048 \
--gpu-memory-utilization 0.8 \
--speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
--max_model_len 2048 \
--gpu_memory_utilization 0.8 \
--speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
```
## Pre-trained PARD weights
@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
```bash
DOCKER_BUILDKIT=1 docker build \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
--build-arg CUDA_VERSION=13.0.1 \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
--build-arg max_jobs=256 \
--build-arg nvcc_threads=2 \
--build-arg RUN_WHEEL_CHECK=false \
@@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt
!!! note
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- Finally, build and install vLLM XPU backend:
+21 -27
View File
@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
|-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
| `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier |
| `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity |
| `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A |
| `token_classify` | Token-wise | probability vector of classes for each token | nan | nan |
| `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) |
!!! note
@@ -86,15 +86,14 @@ The scoring models is designed to compute similarity scores between two input pr
### Pooling Usages
| Pooling Usages | Description |
|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
| Token Classification Usages | Token-wise classification |
| Token Embedding Usages | Token-wise embedding |
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
| Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
| Pooling Usages | Description |
|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
| Token Classification Usages | Token-wise classification |
| Token Embedding Usages | Token-wise embedding |
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.
@@ -102,9 +101,9 @@ For more detailed information, please refer to the link below.
- [Classification Usages](classify.md)
- [Embedding Usages](embed.md)
- [Reward Usages](reward.md)
- [Token Classification Usages](token_classify.md)
- [Token Embedding Usages](token_embed.md)
- [Reward Usages](reward.md)
- [Scoring Usages](scoring.md)
- [Specific Model Examples](specific_models.md)
@@ -114,17 +113,15 @@ Each pooling model in vLLM supports one or more of these tasks according to
[Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
enabling the corresponding APIs.
### Offline APIs corresponding to pooling usages
### Offline APIs corresponding to pooling tasks
| Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function |
|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
| Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier |
| Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity |
| Token Classification Usages | N/A | `token_classify` | N/A | N/A |
| Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) |
| Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A |
| Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A |
| Plugins Usages | N/A | `plugin` | N/A | N/A |
| Task | APIs |
|------------------|---------------------------------------------------------------------------------------|
| `embed` | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)` |
| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` |
| `token_embed` | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)` |
| `plugin` | `LLM.encode(..., pooling_task="plugin")` |
!!! note
Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
@@ -150,7 +147,7 @@ It is primarily designed for [score models](scoring.md).
The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
### Examples
@@ -186,12 +183,9 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all
The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
Code examples:
- [Online example](../../../examples/pooling/reward/token_reward_online.py)
- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)
Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
### Examples
-10
View File
@@ -134,13 +134,3 @@ print(f"Data: {data!r}")
## Online Serving
Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
## More examples
More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
## Deprecated Features
### `LLM.reward`
`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.
+3 -7
View File
@@ -384,7 +384,6 @@ th {
| `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
| `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -420,7 +419,6 @@ th {
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
@@ -474,7 +472,6 @@ th {
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
@@ -563,7 +560,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
| `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + I<sup>E+</sup> | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ |
| `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
| `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |
@@ -644,10 +640,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
!!! note
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
MobileNet-v5 vision backbone.
Performance is not yet fully optimized mainly due to:
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
!!! note
@@ -153,7 +153,6 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
| `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
| `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
| `policy` | The policy type for expert parallel load balancing | `"default"` |
| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` |
For example:
@@ -1,15 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example offline usage of token reward models.
The key distinction between sequence classification and token classification
lies in their output granularity: sequence classification produces a single
result for an entire input sequence, whereas token classification yields a
result for each individual token within the sequence.
"""
from argparse import Namespace
from vllm import LLM, EngineArgs
@@ -45,14 +36,14 @@ def main(args: Namespace):
llm = LLM(**vars(args))
# Generate rewards. The output is a list of PoolingRequestOutput.
outputs = llm.encode(prompts, pooling_task="token_classify")
outputs = llm.reward(prompts)
# Print the outputs.
print("\nGenerated Outputs:\n" + "-" * 60)
for prompt, output in zip(prompts, outputs):
rewards = output.outputs.data
print(f"Prompt: {prompt!r}")
print_embeddings(rewards.tolist(), prefix="Reward")
print_embeddings(rewards, prefix="Reward")
print("-" * 60)
+106 -326
View File
@@ -394,24 +394,18 @@ def run_eagle2_5(questions: list[str], modality: str) -> ModelRequestData:
def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=5,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
trust_remote_code=True,
)
image_placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"
video_placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>"
prompts = [
(
@@ -431,7 +425,6 @@ def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData:
model_name = "LGAI-EXAONE/EXAONE-4.5-33B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -441,23 +434,18 @@ def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData:
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<vision><|image_pad|></vision>"
video_placeholder = "<vision><|video_pad|></vision>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|system|>\nYou are a helpful assistant.<|endofturn|>\n"
f"<|user|>\n{placeholder}"
f"<|user|>\n<vision>{placeholder}</vision>"
f"{question}<|endofturn|>\n"
"<|assistant|>\n"
)
@@ -578,7 +566,6 @@ def run_glm4v(questions: list[str], modality: str) -> ModelRequestData:
def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
model_name = "zai-org/GLM-4.1V-9B-Thinking"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -587,19 +574,14 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
"size": {"shortest_edge": 12544, "longest_edge": 47040000},
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
)
image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
prompts = [
(
@@ -620,7 +602,6 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
model_name = "zai-org/GLM-4.5V"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -629,20 +610,15 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
"size": {"shortest_edge": 12544, "longest_edge": 47040000},
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
tensor_parallel_size=4,
)
image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
prompts = [
(
@@ -663,7 +639,6 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData:
def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
model_name = "zai-org/GLM-4.5V-FP8"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -672,20 +647,15 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
"size": {"shortest_edge": 12544, "longest_edge": 47040000},
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
tensor_parallel_size=4,
)
image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
prompts = [
(
@@ -706,7 +676,6 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
model_name = "zai-org/GLM-OCR"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -715,19 +684,14 @@ def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
"size": {"shortest_edge": 12544, "longest_edge": 47040000},
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
)
image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
prompts = [
(
@@ -808,12 +772,11 @@ def run_hyperclovax_seed_vision(
model_name = "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=16384 if modality in ("video", "image+video") else 8192,
limit_mm_per_prompt=mm_limit,
max_model_len=8192 if modality == "image" else 16384,
limit_mm_per_prompt={modality: 1},
)
messages = list()
@@ -865,29 +828,6 @@ def run_hyperclovax_seed_vision(
}
]
)
elif modality == "image+video":
messages.append(
[
{
"role": "user",
"content": [
{
"type": "image",
"ocr": "",
"lens_keywords": "",
"lens_local_keywords": "",
},
{
"type": "video",
},
{
"type": "text",
"text": question,
},
],
}
]
)
else:
raise ValueError(f"Unsupported modality: {modality}")
@@ -936,25 +876,19 @@ def run_idefics3(questions: list[str], modality: str) -> ModelRequestData:
def run_interns1(questions: list[str], modality: str) -> ModelRequestData:
model_name = "internlm/Intern-S1-mini"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=8192,
max_num_seqs=2,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
)
image_placeholder = "<IMG_CONTEXT>"
video_placeholder = "<video>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<IMG_CONTEXT>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + "\n" + video_placeholder
placeholder = "<video>"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
messages = [
@@ -975,26 +909,20 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData:
def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData:
model_name = "internlm/Intern-S1-Pro"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=8192,
max_num_seqs=2,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
enforce_eager=True,
tensor_parallel_size=4,
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
messages = [
@@ -1015,23 +943,17 @@ def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData:
def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "OpenGVLab/InternVL3-2B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=8192,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<image>"
video_placeholder = "<video>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<image>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + "\n" + video_placeholder
placeholder = "<video>"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
messages = [
@@ -1088,27 +1010,21 @@ def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData:
def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Kwai-Keye/Keye-VL-8B-Preview"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=8192,
trust_remote_code=True,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -1125,27 +1041,21 @@ def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_keye_vl1_5(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Kwai-Keye/Keye-VL-1.5-8B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=8192,
trust_remote_code=True,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -1349,26 +1259,22 @@ def run_llava_next_video(questions: list[str], modality: str) -> ModelRequestDat
# LLaVA-OneVision
def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData:
image_placeholder = "<image>"
video_placeholder = "<video>"
if modality == "video":
prompts = [
f"<|im_start|>user <video>\n{question}<|im_end|><|im_start|>assistant\n"
for question in questions
]
if modality == "image":
placeholder = image_placeholder
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + "\n" + video_placeholder
elif modality == "image":
prompts = [
f"<|im_start|>user <image>\n{question}<|im_end|><|im_start|>assistant\n"
for question in questions
]
prompts = [
(f"<|im_start|>user {placeholder}\n{question}<|im_end|><|im_start|>assistant\n")
for question in questions
]
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model="llava-hf/llava-onevision-qwen2-7b-ov-hf",
max_model_len=16384,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
return ModelRequestData(
@@ -1401,7 +1307,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
# MiniCPM-V
def run_minicpmv_base(questions: list[str], modality: str, model_name):
assert modality in ["image", "video", "image+video"]
assert modality in ["image", "video"]
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
# 2.0
@@ -1423,13 +1329,12 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
# o2.6: image, video, audio
# model_name = "openbmb/MiniCPM-o-2_6"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=2,
trust_remote_code=True,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
# NOTE The stop_token_ids are different for various versions of MiniCPM-V
# 2.0
@@ -1442,22 +1347,17 @@ def run_minicpmv_base(questions: list[str], modality: str, model_name):
stop_tokens = ["<|im_end|>", "<|endoftext|>"]
stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
image_placeholder = "(<image>./</image>)"
video_placeholder = "(<video>./</video>)"
if modality == "image":
placeholder = image_placeholder
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + "\n" + video_placeholder
modality_placeholder = {
"image": "(<image>./</image>)",
"video": "(<video>./</video>)",
}
prompts = [
tokenizer.apply_chat_template(
[
{
"role": "user",
"content": f"{placeholder}\n{question}",
"content": f"{modality_placeholder[modality]}\n{question}",
}
],
tokenize=False,
@@ -1566,24 +1466,20 @@ def run_molmo(questions: list[str], modality: str) -> ModelRequestData:
def run_molmo2(questions: list[str], modality: str) -> ModelRequestData:
model_name = "allenai/Molmo2-8B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
dtype="bfloat16",
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
max_num_batched_tokens=36864,
)
image_placeholder = "<|image|>"
video_placeholder = "<|video|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video|>"
else:
raise ValueError(f"Unsupported modality for molmo2: {modality}")
prompts = [
f"{placeholder}<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1667,25 +1563,19 @@ def run_nvlm_d(questions: list[str], modality: str) -> ModelRequestData:
def run_openpangu_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "FreedomIntelligence/openPangu-VL-7B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=4,
trust_remote_code=True,
enforce_eager=True,
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "[unused19]"
video_placeholder = "[unused32]"
if modality == "image":
placeholder = image_placeholder
placeholder = "[unused19]"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "[unused32]"
prompts = [
(
@@ -1733,25 +1623,18 @@ def run_ovis(questions: list[str], modality: str) -> ModelRequestData:
def run_ovis2_5(questions: list[str], modality: str) -> ModelRequestData:
model_name = "AIDC-AI/Ovis2.5-2B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=2,
trust_remote_code=True,
dtype="half",
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<image>"
video_placeholder = "<video>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<image>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + "\n" + video_placeholder
placeholder = "<video>"
prompts = [
f"<|im_start|>user\n\n{placeholder}\n{question}<|im_end|>\n<|im_start|>assistant\n"
@@ -1963,7 +1846,6 @@ def run_qwen_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen2-VL-7B-Instruct"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -1973,23 +1855,18 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
"min_pixels": 28 * 28,
"max_pixels": 1280 * 28 * 28,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2006,7 +1883,6 @@ def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen2.5-VL-3B-Instruct"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -2016,23 +1892,18 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2049,7 +1920,6 @@ def run_qwen2_5_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_qwen2_5_omni(questions: list[str], modality: str):
model_name = "Qwen/Qwen2.5-Omni-7B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -2059,18 +1929,13 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_bos|><|IMAGE|><|vision_eos|>"
video_placeholder = "<|vision_bos|><|VIDEO|><|vision_eos|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|IMAGE|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|VIDEO|>"
default_system = (
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
@@ -2081,7 +1946,7 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
prompts = [
(
f"<|im_start|>system\n{default_system}<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_bos|>{placeholder}<|vision_eos|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2097,7 +1962,6 @@ def run_qwen2_5_omni(questions: list[str], modality: str):
def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen3-VL-4B-Instruct"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -2107,23 +1971,18 @@ def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2140,7 +1999,6 @@ def run_qwen3_vl(questions: list[str], modality: str) -> ModelRequestData:
def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen3-VL-30B-A3B-Instruct"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -2150,23 +2008,18 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2337,7 +2190,6 @@ def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
model_name = "omni-research/Tarsier2-Recap-7b"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
@@ -2345,23 +2197,18 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
"architectures": ["Tarsier2ForConditionalGeneration"],
"model_type": "tarsier2",
},
limit_mm_per_prompt=mm_limit,
limit_mm_per_prompt={modality: 1},
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
placeholder = "<|image_pad|>"
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
placeholder = "<|video_pad|>"
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
@@ -2463,12 +2310,6 @@ MODELS_NEED_VIDEO_METADATA = [
]
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
"qwen3_vl",
"qwen3_vl_moe",
]
def get_multi_modal_input(args):
"""
return {
@@ -2516,24 +2357,6 @@ def get_multi_modal_input(args):
"questions": vision_chunk_questions,
}
if args.modality == "image+video":
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
needs_metadata = args.model_type in MODELS_NEED_VIDEO_METADATA
video = VideoAsset(name="baby_reading", num_frames=args.num_frames).np_ndarrays
metadata = VideoAsset(name="baby_reading", num_frames=args.num_frames).metadata
img_video_questions = [
"What is shown in the image? What happens in the video?",
"Describe both the image and the video content.",
]
return {
"data": {
"image": image,
"video": ([(video, metadata)] if needs_metadata else video),
},
"questions": img_video_questions,
}
msg = f"Modality {args.modality} is not supported."
raise ValueError(msg)
@@ -2581,29 +2404,6 @@ def apply_image_repeat(
return inputs, inputs_with_empty_media
def maybe_add_vit_cuda_graph_compilation_config(args, engine_args):
model = args.model_type
modality = args.modality
enable_vit_cuda_graph = args.enable_vit_cuda_graph
if enable_vit_cuda_graph and model in MODELS_SUPPORT_VIT_CUDA_GRAPH:
if modality == "image" or modality == "video":
vision_items_per_batch = 1
elif modality == "image+video":
vision_items_per_batch = 2
else:
raise ValueError(
f"modality={modality} is not supported for vit cuda graph."
)
engine_args.compilation_config = {
"cudagraph_mm_encoder": True,
"encoder_cudagraph_max_vision_items_per_batch": vision_items_per_batch,
}
return engine_args
@contextmanager
def time_counter(enable: bool):
if enable:
@@ -2639,7 +2439,7 @@ def parse_args():
"--modality",
type=str,
default="image",
choices=["image", "video", "image+video", "vision_chunk"],
choices=["image", "video", "vision_chunk"],
help="Modality of the input.",
)
parser.add_argument(
@@ -2654,28 +2454,33 @@ def parse_args():
default=0,
help="Set the seed when initializing `vllm.LLM`.",
)
parser.add_argument(
"--image-repeat-prob",
type=float,
default=None,
help="Simulates the hit-ratio for multi-modal preprocessor cache (if enabled)",
)
parser.add_argument(
"--disable-mm-processor-cache",
action="store_true",
help="If True, disables caching of multi-modal processor.",
)
parser.add_argument(
"--time-generate",
action="store_true",
help="If True, then print the total generate() call time",
)
parser.add_argument(
"--use-different-prompt-per-request",
action="store_true",
help="If True, then use different prompt (with the same multi-modal "
"data) for each request.",
)
parser.add_argument(
"--verify-mm-cache-hit-with-uuids",
action="store_true",
@@ -2689,11 +2494,6 @@ def parse_args():
default=None,
help="Tensor parallel size to override the model's default setting. ",
)
parser.add_argument(
"--enable-vit-cuda-graph",
action="store_true",
help="If True, will enable vit cuda graph capture and replay for the model.",
)
return parser.parse_args()
@@ -2727,7 +2527,6 @@ def main(args):
engine_args.mm_processor_cache_gb = mm_processor_cache_gb
if args.tensor_parallel_size is not None:
engine_args.tensor_parallel_size = args.tensor_parallel_size
engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
llm = LLM.from_engine_args(engine_args)
# Don't want to check the flag multiple times, so just hijack `prompts`.
@@ -2747,42 +2546,23 @@ def main(args):
else req_data.sampling_params
)
def _mm_data(data, modality):
if modality == "image+video":
return {"image": data["image"], "video": data["video"]}
return {modality: data}
def _mm_uuid(uuid, modality):
if modality == "image+video":
return {"image": uuid, "video": uuid + "v"}
return {modality: uuid}
def _mm_empty(modality):
if modality == "image+video":
return {"image": None, "video": None}
return {modality: None}
assert args.num_prompts > 0
if args.num_prompts == 1:
# Single inference
uuid = "uuid_0"
inputs = {
"prompt": prompts[0],
"multi_modal_data": _mm_data(data, modality),
"multi_modal_uuids": _mm_uuid(uuid, modality),
"multi_modal_data": {modality: data},
"multi_modal_uuids": {modality: uuid},
}
inputs_with_empty_media = {
"prompt": prompts[0],
"multi_modal_data": _mm_empty(modality),
"multi_modal_uuids": _mm_uuid(uuid, modality),
"multi_modal_data": {modality: None},
"multi_modal_uuids": {modality: uuid},
}
else:
# Batch inference
if args.image_repeat_prob is not None:
if modality == "image+video":
raise ValueError(
"--image-repeat-prob is not supported for 'image+video' modality"
)
# Repeat images with specified probability of "image_repeat_prob"
inputs, inputs_with_empty_media = apply_image_repeat(
args.image_repeat_prob,
@@ -2792,7 +2572,7 @@ def main(args):
modality,
)
else:
# Use the same image/video for all prompts
# Use the same image for all prompts
inputs = []
inputs_with_empty_media = []
for i in range(args.num_prompts):
@@ -2800,15 +2580,15 @@ def main(args):
inputs.append(
{
"prompt": prompts[i % len(prompts)],
"multi_modal_data": _mm_data(data, modality),
"multi_modal_uuids": _mm_uuid(uuid, modality),
"multi_modal_data": {modality: data},
"multi_modal_uuids": {modality: uuid},
}
)
inputs_with_empty_media.append(
{
"prompt": prompts[i % len(prompts)],
"multi_modal_data": _mm_empty(modality),
"multi_modal_uuids": _mm_uuid(uuid, modality),
"multi_modal_data": {modality: None},
"multi_modal_uuids": {modality: uuid},
}
)
@@ -310,38 +310,6 @@ def load_gemma3(question: str, image_urls: list[str]) -> ModelRequestData:
)
def load_granite4_vision(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "ibm-granite/granite-vision-4.1-4b"
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=16,
limit_mm_per_prompt={"image": len(image_urls)},
)
placeholders = [{"type": "image", "image": url} for url in image_urls]
messages = [
{
"role": "user",
"content": [
*placeholders,
{"type": "text", "text": question},
],
}
]
processor = AutoProcessor.from_pretrained(model_name)
prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
image_data=[fetch_image(url) for url in image_urls],
)
def load_h2ovl(question: str, image_urls: list[str]) -> ModelRequestData:
model_name = "h2oai/h2ovl-mississippi-800m"
@@ -1519,7 +1487,6 @@ model_example_map = {
"deepseek_ocr": load_deepseek_ocr,
"exaone4_5": load_exaone4_5,
"gemma3": load_gemma3,
"granite4_vision": load_granite4_vision,
"h2ovl_chat": load_h2ovl,
"hunyuan_vl": load_hunyuan_vl,
"hyperclovax_seed_vision": load_hyperclovax_seed_vision,
@@ -10,8 +10,9 @@ import uuid
import aiohttp
import msgpack
import regex as re
import zmq
from quart import Quart, Request, make_response, request
from quart import Quart, make_response, request
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
MoRIIOConstants,
@@ -24,10 +25,32 @@ decode_instances: list[dict] = []
request_nums = 0
app = Quart(__name__)
IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
TRANSFER_TYPE = None
def _append_whole_dict_unique(target_list, data_dict):
new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
for existed in target_list:
existed_filtered = {k: v for k, v in existed.items() if k != "index"}
if existed_filtered == new_filtered:
return False
print("!!APPEND!!", data_dict)
target_list.append(data_dict)
transfer_mode = data_dict.get("transfer_mode", "unknown")
global TRANSFER_TYPE
if TRANSFER_TYPE is None:
TRANSFER_TYPE = transfer_mode
logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
elif transfer_mode != TRANSFER_TYPE:
raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
return True
_list_lock = threading.RLock()
@@ -45,81 +68,23 @@ def _listen_for_register(hostname, port):
if router_socket in socks:
remote_addr, msg = router_socket.recv_multipart()
data = msgpack.loads(msg)
if data.get("type") == "HELLO":
if data["type"] == "HELLO":
pass
elif data.get("type") in ("P", "D"):
role = data["type"]
required_keys = {
"http_address",
"zmq_address",
"dp_size",
"tp_size",
"transfer_mode",
}
missing = required_keys - data.keys()
if missing:
logger.error(
"Registration message missing required keys %s; skipping",
missing,
)
continue
# Derive request_address from http_address
# api path suffix is appended at request time
instance = {
"role": role,
"request_address": f"http://{data['http_address']}/v1",
"http_address": data["http_address"],
"zmq_address": data["zmq_address"],
"dp_size": data["dp_size"],
"tp_size": data["tp_size"],
"transfer_mode": data["transfer_mode"],
}
# zmq_address format: "host:IP,handshake:PORT,notify:PORT"
# Stored verbatim; embedded into the request_id by handle_request.
global TRANSFER_TYPE
transfer_mode = instance["transfer_mode"]
target_list = prefill_instances if role == "P" else decode_instances
elif (
data["type"] == "register"
and data["role"] == "P"
and data["request_address"] not in prefill_instances
):
with _list_lock:
if TRANSFER_TYPE is None:
TRANSFER_TYPE = transfer_mode
logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
elif transfer_mode != TRANSFER_TYPE:
logger.error(
"Mismatched transfer mode: expected %s, got %s;"
" skipping registration of %s",
TRANSFER_TYPE,
transfer_mode,
data["http_address"],
)
continue
existing_idx = next(
(
idx
for idx, i in enumerate(target_list)
if i.get("http_address") == data["http_address"]
),
None,
)
if existing_idx is not None:
target_list[existing_idx] = instance
logger.info(
"Updated existing %s instance: %s",
"Prefill" if role == "P" else "Decode",
instance,
)
else:
target_list.append(instance)
logger.info(
"Registered %s instance: %s",
"Prefill" if role == "P" else "Decode",
instance,
)
else:
logger.warning(
"Received message with unrecognized type %r; ignoring",
data.get("type"),
)
_append_whole_dict_unique(prefill_instances, data)
elif (
data["type"] == "register"
and data["role"] == "D"
and data["request_address"] not in decode_instances
):
with _list_lock:
_append_whole_dict_unique(decode_instances, data)
def start_service_discovery(hostname, port):
@@ -136,7 +101,7 @@ def start_service_discovery(hostname, port):
async def send_request_to_prefill(
endpoint, req_data, request_id, selected_prefill_dp_rank
endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
):
req_data_copy = req_data
@@ -144,8 +109,12 @@ async def send_request_to_prefill(
{
"do_remote_decode": True,
"do_remote_prefill": False,
"remote_handshake_port": d_endpoint["handshake_port"],
"remote_notify_port": d_endpoint["notify_port"],
"remote_engine_id": None,
"remote_block_ids": None,
"remote_host": dip,
"remote_port": dport,
}
)
req_data_copy["stream"] = False
@@ -170,13 +139,10 @@ async def send_request_to_prefill(
return await response.json()
else:
error_message = (
f"send_request_to_prefill response ={response},"
f"reason={response.reason}, status={response.status},"
f"method={response.method}, url={response.url},"
f"real_url={response.real_url}"
raise RuntimeError(
"send_request_to_prefill response.status != 200response.status = ",
response.status,
)
raise RuntimeError(error_message)
async def start_decode_request(endpoint, req_data, request_id):
@@ -197,13 +163,9 @@ async def stream_decode_response(session, response, request_id):
async for chunk_bytes in response.content.iter_chunked(1024):
yield chunk_bytes
else:
error_message = (
f"stream_decode_response response ={response},"
f"reason={response.reason}, status={response.status},"
f"method={response.method}, url={response.url},"
f"real_url={response.real_url}"
raise RuntimeError(
f"decode response.status != 200, status = {response.status}"
)
raise RuntimeError(error_message)
finally:
await session.close()
@@ -213,22 +175,21 @@ def example_round_robin_dp_loader(request_number, dp_size):
@app.route("/v1/completions", methods=["POST"])
async def handle_completions_request():
return await handle_request("/completions", request)
@app.route("/v1/chat/completions", methods=["POST"])
async def handle_chat_completions_request():
return await handle_request("/chat/completions", request)
async def handle_request(api: str, request: Request):
async def handle_request():
try:
with _list_lock:
global request_nums
request_nums += 1
def extract_ip_port_fast(url):
match = IP_PORT_PATTERN.search(url)
if not match:
raise ValueError(f"Invalid URL format: {url}")
return match.groups()
req_data = await request.get_json()
request_id = str(uuid.uuid4())
prefill_instance_endpoint = None
decode_instance_endpoint = None
@@ -254,14 +215,7 @@ async def handle_request(api: str, request: Request):
prefill_instance_endpoint["dp_size"],
)
# Embed both zmq_addresses in the request_id so the connector can parse
# the peer's host/ports from it, similar to P2P-NCCL
uid = str(uuid.uuid4()).replace("-", "")
request_id = (
f"___prefill_addr_{prefill_instance_endpoint['zmq_address']}"
f"___decode_addr_{decode_instance_endpoint['zmq_address']}"
f"_{uid}"
)
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
@@ -276,36 +230,40 @@ async def handle_request(api: str, request: Request):
)
req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
prefill_request_url = prefill_instance_endpoint["request_address"] + api
send_prefill_task = asyncio.create_task(
send_request_to_prefill(
prefill_request_url,
prefill_instance_endpoint["request_address"],
req_data_to_prefill,
request_id,
decode_instance_endpoint,
dip,
dport,
selected_prefill_dp_rank,
)
)
ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
req_data["max_tokens"] -= 1
req_data["kv_transfer_params"] = {
"do_remote_decode": False,
"do_remote_prefill": True,
"remote_handshake_port": prefill_instance_endpoint["handshake_port"],
"remote_notify_port": prefill_instance_endpoint["notify_port"],
"remote_engine_id": None,
"remote_block_ids": None,
"transfer_id": transfer_id,
"remote_host": ip,
"remote_port": port,
}
if TRANSFER_TYPE == "READ":
# In read mode, prefill and decode are executed serially.
prefill_response = await send_prefill_task
prefill_kv = prefill_response["kv_transfer_params"]
req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[
"remote_engine_id"
]
req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[
"remote_block_ids"
]
req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
"kv_transfer_params"
]["remote_engine_id"]
req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
"kv_transfer_params"
]["remote_block_ids"]
req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
"dp_size"
@@ -316,10 +274,12 @@ async def handle_request(api: str, request: Request):
if selected_prefill_dp_rank is not None:
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
req_data["kv_transfer_params"]["transfer_id"] = transfer_id
decode_request_url = decode_instance_endpoint["request_address"] + api
decode_request_task = asyncio.create_task(
start_decode_request(decode_request_url, req_data, request_id)
start_decode_request(
decode_instance_endpoint["request_address"], req_data, request_id
)
)
session, decode_response = await decode_request_task
@@ -27,12 +27,7 @@ from vllm.assets.audio import AudioAsset
def sync_openai(
audio_path: str,
client: OpenAI,
model: str,
*,
repetition_penalty: float = 1.3,
hotwords: str = None,
audio_path: str, client: OpenAI, model: str, *, repetition_penalty: float = 1.3
):
"""
Perform synchronous transcription using OpenAI-compatible API.
@@ -48,15 +43,12 @@ def sync_openai(
extra_body=dict(
seed=4419,
repetition_penalty=repetition_penalty,
hotwords=hotwords,
),
)
print("transcription result [sync]:", transcription.text)
async def stream_openai_response(
audio_path: str, client: AsyncOpenAI, model: str, hotwords: str = None
):
async def stream_openai_response(audio_path: str, client: AsyncOpenAI, model: str):
"""
Perform asynchronous transcription using OpenAI-compatible API.
"""
@@ -72,7 +64,6 @@ async def stream_openai_response(
extra_body=dict(
seed=420,
top_p=0.6,
hotwords=hotwords,
),
stream=True,
)
@@ -145,7 +136,6 @@ def main(args):
client=client,
model=model,
repetition_penalty=args.repetition_penalty,
hotwords=args.hotwords,
)
# Run the asynchronous function
@@ -156,10 +146,7 @@ def main(args):
)
asyncio.run(
stream_openai_response(
args.audio_path if args.audio_path else winning_call,
client,
model,
hotwords=args.hotwords,
args.audio_path if args.audio_path else winning_call, client, model
)
)
else:
@@ -187,11 +174,5 @@ if __name__ == "__main__":
default=1.3,
help="repetition penalty",
)
parser.add_argument(
"--hotwords",
type=str,
default=None,
help="hotwords",
)
args = parser.parse_args()
main(args)
@@ -1,17 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example online usage of token reward models.
Example online usage of Pooling API.
Run `vllm serve <model> --runner pooling`
to start up the server in vLLM. e.g.
vllm serve internlm/internlm2-1_8b-reward --trust-remote-code
The key distinction between sequence classification and token classification
lies in their output granularity: sequence classification produces a single
result for an entire input sequence, whereas token classification yields a
result for each individual token within the sequence.
"""
import argparse
@@ -1,62 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example offline usage of sequence reward models.
The key distinction between sequence classification and token classification
lies in their output granularity: sequence classification produces a single
result for an entire input sequence, whereas token classification yields a
result for each individual token within the sequence.
"""
from argparse import Namespace
from vllm import LLM, EngineArgs
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.utils.print_utils import print_embeddings
def parse_args():
parser = FlexibleArgumentParser()
parser = EngineArgs.add_cli_args(parser)
# Set example specific arguments
parser.set_defaults(
model="Skywork/Skywork-Reward-V2-Qwen3-0.6B",
runner="pooling",
enforce_eager=True,
max_model_len=1024,
trust_remote_code=True,
)
return parser.parse_args()
def main(args: Namespace):
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create an LLM.
# You should pass runner="pooling" for reward models
llm = LLM(**vars(args))
# Generate rewards. The output is a list of PoolingRequestOutput.
# Use pooling_task="classify" for sequence reward models.
outputs = llm.encode(prompts, pooling_task="classify")
# Print the outputs.
print("\nGenerated Outputs:\n" + "-" * 60)
for prompt, output in zip(prompts, outputs):
rewards = output.outputs.data
print(f"Prompt: {prompt!r}")
print_embeddings(rewards.tolist(), prefix="Reward")
print("-" * 60)
if __name__ == "__main__":
args = parse_args()
main(args)
@@ -1,71 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example online usage of sequence reward models.
Run `vllm serve <model> --runner pooling`
to start up the server in vLLM. e.g.
vllm serve Skywork/Skywork-Reward-V2-Qwen3-0.6B
The key distinction between sequence classification and token classification
lies in their output granularity: sequence classification produces a single
result for an entire input sequence, whereas token classification yields a
result for each individual token within the sequence.
"""
import argparse
import pprint
import requests
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
headers = {"User-Agent": "Test Client"}
response = requests.post(api_url, headers=headers, json=prompt)
return response
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--port", type=int, default=8000)
return parser.parse_args()
def main(args):
base_url = f"http://{args.host}:{args.port}"
models_url = base_url + "/v1/models"
pooing_url = base_url + "/pooling"
response = requests.get(models_url)
model = response.json()["data"][0]["id"]
# Input like Completions API
prompt = {"model": model, "input": "vLLM is great!"}
pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
print("-" * 50)
print("Pooling Response:")
pprint.pprint(pooling_response.json())
print("-" * 50)
# Input like Chat API
prompt = {
"model": model,
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "vLLM is great!"}],
}
],
}
pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
print("Pooling Response:")
pprint.pprint(pooling_response.json())
print("-" * 50)
if __name__ == "__main__":
args = parse_args()
main(args)
+8 -1
View File
@@ -131,9 +131,16 @@ class TrainModel:
from vllm.model_executor.layers.batch_invariant import (
init_batch_invariance,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backends.registry import AttentionBackendEnum
# need to init all env vars for batch invariance which affect nccl ops
init_batch_invariance()
attn_backend = (
AttentionBackendEnum.TRITON_ATTN
if current_platform.is_rocm()
else AttentionBackendEnum.FLASH_ATTN
)
init_batch_invariance(attn_backend)
self.model = AutoModelForCausalLM.from_pretrained(
model_name, dtype=torch.bfloat16

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