Compare commits

..
Author SHA1 Message Date
Wentao YeandGitHub beefff2174 Merge branch 'main' into wentao-skip-work-when-empty 2026-04-03 11:31:23 -04:00
yewentao256 163266d0b2 refactor only inside nixl
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-01 18:28:13 +00:00
yewentao256 a2fd28a7e1 Merge branch 'main' into wentao-skip-work-when-empty
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-01 18:09:06 +00:00
Wentao YeandGitHub 7a80ac928f Merge branch 'main' into wentao-skip-work-when-empty 2026-03-31 13:01:59 -04:00
yewentao256 0ed11013b4 check empty using kv connector's method
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-03-30 21:06:12 +00:00
yewentao256 6d568b995a Merge branch 'main' into wentao-skip-work-when-empty 2026-03-30 19:17:12 +00:00
Wentao YeandGitHub dfe9decbcb Merge branch 'main' into wentao-skip-work-when-empty 2026-03-28 10:28:43 -04:00
Wentao YeandGitHub 1324e6ff67 Merge branch 'main' into wentao-skip-work-when-empty 2026-03-27 15:52:55 -04:00
Wentao YeandGitHub c1aba6d7ae Merge branch 'main' into wentao-skip-work-when-empty 2026-03-27 09:42:19 -04:00
Wentao YeGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
200bef28c9 Update vllm/v1/worker/gpu/kv_connector.py
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-03-26 18:04:12 -04:00
yewentao256 fd9820bbf9 skip kv connector empty work
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-03-26 21:50:28 +00:00
1240 changed files with 30618 additions and 94701 deletions
+2 -2
View File
@@ -8,8 +8,8 @@ run_all_patterns:
- "CMakeLists.txt"
- "requirements/common.txt"
- "requirements/cuda.txt"
- "requirements/build/cuda.txt"
- "requirements/test/cuda.txt"
- "requirements/build.txt"
- "requirements/test.txt"
- "setup.py"
- "csrc/"
- "cmake/"
+2 -2
View File
@@ -6,8 +6,8 @@ run_all_patterns:
- "CMakeLists.txt"
- "requirements/common.txt"
- "requirements/xpu.txt"
- "requirements/build/cuda.txt"
- "requirements/test/cuda.txt"
- "requirements/build.txt"
- "requirements/test.txt"
- "setup.py"
- "csrc/"
- "cmake/"
+1
View File
@@ -5,6 +5,7 @@ steps:
depends_on: []
device: amd_cpu
no_plugin: true
soft_fail: true
commands:
- >
docker build
+2 -2
View File
@@ -46,7 +46,7 @@ steps:
- tests/models/language/pooling/
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
pytest -x -v -s tests/models/language/generation -m cpu_model
pytest -x -v -s tests/models/language/pooling -m cpu_model"
@@ -99,7 +99,7 @@ steps:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
parallelism: 3
parallelism: 2
- label: "Arm CPU Test"
depends_on: []
+2 -2
View File
@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
}
ecr_login() {
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
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
}
prepare_cache_tags() {
+1 -1
View File
@@ -11,7 +11,7 @@ REPO=$2
BUILDKITE_COMMIT=$3
# authenticate with AWS ECR
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
# skip build if image already exists
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
@@ -11,7 +11,7 @@ REPO=$2
BUILDKITE_COMMIT=$3
# authenticate with AWS ECR
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
# skip build if image already exists
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
@@ -1,68 +0,0 @@
#!/bin/bash
set -euo pipefail
# Build a vLLM test image with PyTorch nightly installed.
# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
if [[ $# -lt 5 ]]; then
echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
exit 1
fi
REGISTRY=$1
REPO=$2
BUILDKITE_COMMIT=$3
BRANCH=$4
IMAGE_TAG=$5
# --- Arguments ---
echo "--- :mag: Arguments"
echo "REGISTRY: ${REGISTRY}"
echo "REPO: ${REPO}"
echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
echo "BRANCH: ${BRANCH}"
echo "IMAGE_TAG: ${IMAGE_TAG}"
# --- ECR login ---
echo "--- :key: ECR login"
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
# --- Set up buildx ---
echo "--- :docker: Setting up buildx"
docker buildx create --name vllm-builder --driver docker-container --use || true
docker buildx inspect --bootstrap
docker buildx ls
# --- Skip if image already exists ---
echo "--- :mag: Checking if image already exists"
if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
echo "Image found: $IMAGE_TAG — skipping build"
exit 0
fi
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_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
docker buildx build --file docker/Dockerfile \
--build-arg max_jobs=16 \
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
--build-arg USE_SCCACHE=1 \
--build-arg PYTORCH_NIGHTLY=1 \
--build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
--build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
--build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
--build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
--tag "$IMAGE_TAG" \
--push \
--target test \
--progress plain .
echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"
-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'
+1 -2
View File
@@ -35,7 +35,6 @@ steps:
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel'
@@ -62,4 +61,4 @@ steps:
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py'
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py'
@@ -1,9 +1,6 @@
# For hf script, without -t option (tensor parallel size).
# bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
required_gpu_arch:
- gfx942
- gfx950
tasks:
- name: "mmlu_pro"
metrics:
@@ -1,9 +1,6 @@
# For vllm script, with -t option (tensor parallel size)
# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic -l 1319 -t 1
model_name: "RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic"
required_gpu_arch:
- gfx942
- gfx950
tasks:
- name: "gsm8k"
metrics:
@@ -1,7 +1,4 @@
model_name: "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8"
required_gpu_arch:
- gfx942
- gfx950
tasks:
- name: "mmlu_pro"
metrics:
@@ -1,6 +1,5 @@
Qwen2.5-1.5B-Instruct.yaml
Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
Meta-Llama-3-8B-Instruct-INT8-compressed-tensors-asym.yaml
Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
Qwen1.5-MoE-W4A16-compressed-tensors.yaml
@@ -13,7 +13,6 @@ import os
from contextlib import contextmanager
import lm_eval
import pytest
import yaml
from vllm.platforms import current_platform
@@ -90,40 +89,9 @@ def launch_lm_eval(eval_config, tp_size):
return results
def _check_rocm_gpu_arch_requirement(eval_config):
"""Skip the test if the model requires a ROCm GPU arch not present.
Model YAML configs can specify::
required_gpu_arch:
- gfx942
- gfx950
The check only applies on ROCm. On other platforms (e.g. CUDA) the
field is ignored so that shared config files work for both NVIDIA and
AMD CI pipelines.
"""
required_archs = eval_config.get("required_gpu_arch")
if not required_archs:
return
if not current_platform.is_rocm():
return
from vllm.platforms.rocm import _GCN_ARCH # noqa: E402
if not any(arch in _GCN_ARCH for arch in required_archs):
pytest.skip(
f"Model requires GPU arch {required_archs}, "
f"but detected arch is '{_GCN_ARCH}'"
)
def test_lm_eval_correctness_param(config_filename, tp_size):
eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
_check_rocm_gpu_arch_requirement(eval_config)
results = launch_lm_eval(eval_config, tp_size)
rtol = eval_config.get("rtol", DEFAULT_RTOL)
+57 -71
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"
@@ -104,105 +98,99 @@ steps:
commands:
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
- block: "Unblock to build release Docker images"
depends_on: ~
key: block-build-release-images
if: build.env("NIGHTLY") != "1"
- group: "Build release Docker images"
key: "build-release-images"
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 +231,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 +243,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 +252,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 +276,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 +306,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
@@ -629,8 +617,6 @@ steps:
- label: ":docker: Build release image - x86_64 - ROCm"
id: build-rocm-release-image
depends_on:
- step: block-build-release-images
allow_failure: true
- step: build-rocm-base-wheels
allow_failure: false
agents:
+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/"
@@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12)
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
@@ -35,6 +35,23 @@ export PYTHONPATH=".."
# Helper Functions
###############################################################################
wait_for_clean_gpus() {
local timeout=${1:-300}
local start=$SECONDS
echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
while true; do
if grep -q clean /opt/amdgpu/etc/gpu_state; then
echo "GPUs state is \"clean\""
return
fi
if (( SECONDS - start >= timeout )); then
echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
exit 1
fi
sleep 3
done
}
cleanup_docker() {
# Get Docker's root directory
docker_root=$(docker info -f '{{.DockerRootDir}}')
@@ -348,12 +365,19 @@ apply_rocm_test_overrides() {
###############################################################################
# --- GPU initialization ---
echo "--- Confirming Clean Initial State"
wait_for_clean_gpus
echo "--- ROCm info"
rocminfo
# --- Docker housekeeping ---
cleanup_docker
echo "--- Resetting GPUs"
echo "reset" > /opt/amdgpu/etc/gpu_state
wait_for_clean_gpus
# --- Pull test image ---
echo "--- Pulling container"
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
exit 1
fi
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
#echo "--- DP+TP"
#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
#server_pid=$!
#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
#vllm bench serve \
# --backend vllm \
# --dataset-name random \
# --model meta-llama/Llama-3.2-3B-Instruct \
# --num-prompts 20 \
# --result-dir ./test_results \
# --result-filename dp_pp.json \
# --save-result \
# --endpoint /v1/completions
#kill -s SIGTERM $server_pid; wait $server_pid || true
#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
#if [ "$failed_req" -ne 0 ]; then
# echo "Some requests were failed!"
# exit 1
#fi
@@ -51,7 +51,6 @@ function cpu_tests() {
set -e
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
# basic online serving
@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
# Run the image, setting --shm-size=4g for tensor parallel.
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
@@ -42,7 +42,7 @@ WORKDIR /workspace/vllm
ENV no_proxy=localhost,127.0.0.1
ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true
RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build/cuda.txt)'
RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build.txt)'
RUN VLLM_TARGET_DEVICE=empty pip install --no-build-isolation -e .
RUN pip install git+https://github.com/vllm-project/vllm-gaudi.git
@@ -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
@@ -48,6 +50,6 @@ docker run \
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
pytest -v -s v1/structured_output
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)"
pytest -v -s v1/test_serial_utils.py
'
@@ -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
+2954 -2564
View File
File diff suppressed because it is too large Load Diff
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: Basic Correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/basic_correctness/test_basic_correctness
-1
View File
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: Benchmarks CLI Test
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/benchmarks/
-1
View File
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: Platform Tests (CUDA)
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/cuda
-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
+71 -2
View File
@@ -196,8 +196,6 @@ steps:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- label: Distributed Tests (2 GPUs)(B200)
device: b200
@@ -226,6 +224,77 @@ 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: MessageQueue TCP Multi-Node (2 GPUs)
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
num_devices: 1
num_nodes: 2
no_plugin: true
optional: true
source_file_dependencies:
- vllm/distributed/device_communicators/shm_broadcast.py
- vllm/distributed/parallel_state.py
- tests/distributed/test_mq_tcp_multinode.py
commands:
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 1 $IMAGE_TAG "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py" "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py"
- 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: 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: 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
-2
View File
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: Engine
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/engine
@@ -26,7 +25,6 @@ steps:
- label: e2e Scheduling (1 GPU)
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
-2
View File
@@ -61,7 +61,6 @@ steps:
- label: Entrypoints Integration (API Server openai - Part 3)
timeout_in_minutes: 50
device: h200_18gb
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -106,7 +105,6 @@ steps:
- label: OpenAI API Correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: EPLB Algorithm
timeout_in_minutes: 15
device: h200_18gb
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/eplb
+3 -51
View File
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: vLLM IR Tests
timeout_in_minutes: 10
device: h200_18gb
working_dir: "/vllm-workspace/"
source_file_dependencies:
- vllm/ir
@@ -18,22 +17,10 @@ steps:
source_file_dependencies:
- csrc/
- tests/kernels/core
- tests/kernels/test_top_k_per_row.py
- tests/kernels/test_concat_mla_q.py
commands:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
timeout_in_minutes: 15
num_devices: 2
device: h100
source_file_dependencies:
- csrc/minimax_reduce_rms_kernel.cu
- csrc/minimax_reduce_rms_kernel.h
- vllm/model_executor/layers/mamba/linear_attn.py
- vllm/model_executor/layers/mamba/lamport_workspace.py
- tests/kernels/core/test_minimax_reduce_rms.py
commands:
- pytest -v -s kernels/core/test_minimax_reduce_rms.py
- pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
- label: Kernels Attention Test %N
timeout_in_minutes: 35
@@ -119,7 +106,6 @@ steps:
- vllm/v1/attention/backends/mla/flashinfer_mla.py
- vllm/v1/attention/selector.py
- vllm/platforms/cuda.py
- tests/kernels/test_top_k_per_row.py
commands:
- nvidia-smi
- python3 examples/basic/offline_inference/chat.py
@@ -130,7 +116,6 @@ steps:
- pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
- pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
- pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
- pytest -v -s tests/kernels/test_top_k_per_row.py
# Quantization
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
@@ -141,7 +126,6 @@ steps:
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
- pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
- pytest -v -s tests/kernels/moe/test_flashinfer.py
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
@@ -156,7 +140,7 @@ steps:
- vllm/utils/import_utils.py
- tests/kernels/helion/
commands:
- pip install helion==1.0.0
- pip install helion==0.3.3
- pytest -v -s kernels/helion/
@@ -195,35 +179,3 @@ steps:
- pytest -v -s kernels/moe/test_flashinfer_moe.py
- pytest -v -s kernels/moe/test_nvfp4_moe.py
- pytest -v -s kernels/moe/test_ocp_mx_moe.py
- label: Kernels FusedMoE Layer Test (2 H100s)
timeout_in_minutes: 90
device: h100
num_devices: 2
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
- tests/kernels/moe
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/device_communicators/
- vllm/config
commands:
- pytest -v -s kernels/moe/test_moe_layer.py
- label: Kernels FusedMoE Layer Test (2 B200s)
timeout_in_minutes: 90
device: b200
num_devices: 2
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
- tests/kernels/moe
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/device_communicators/
- vllm/config
commands:
- pytest -v -s kernels/moe/test_moe_layer.py
-10
View File
@@ -91,16 +91,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
- label: LM Eval TurboQuant KV Cache
timeout_in_minutes: 75
source_file_dependencies:
- vllm/model_executor/layers/quantization/turboquant/
- vllm/v1/attention/backends/turboquant_attn.py
- vllm/v1/attention/ops/triton_turboquant_decode.py
- vllm/v1/attention/ops/triton_turboquant_store.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
- label: GPQA Eval (GPT-OSS) (H100)
timeout_in_minutes: 120
device: h100
-5
View File
@@ -19,7 +19,6 @@ steps:
- label: V1 Sample + Logits
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/v1/sample
@@ -87,7 +86,6 @@ steps:
- label: Regression
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/test_regression
@@ -176,7 +174,6 @@ steps:
- tests/renderers
- tests/standalone_tests/lazy_imports.py
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/transformers_utils
- tests/config
@@ -190,7 +187,6 @@ steps:
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s renderers
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s transformers_utils
- pytest -v -s config
@@ -224,7 +220,6 @@ steps:
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- label: Acceptance Length Test (Large Models) # optional
timeout_in_minutes: 25
+1 -2
View File
@@ -78,6 +78,7 @@ steps:
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
# These require fix https://github.com/vllm-project/vllm/pull/36280
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
@@ -100,13 +101,11 @@ steps:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+4 -22
View File
@@ -1,5 +1,5 @@
group: Models - Basic
depends_on:
depends_on:
- image-build
steps:
- label: Basic Models Tests (Initialization)
@@ -12,11 +12,10 @@ steps:
commands:
# Run a subset of model initialization tests
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
mirror:
torch_nightly: {}
- label: Basic Models Tests (Extra Initialization) %N
timeout_in_minutes: 45
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
- tests/models/test_initialization.py
@@ -27,8 +26,6 @@ steps:
# test.) Also run if model initialization test file is modified
- pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
parallelism: 2
mirror:
torch_nightly: {}
- label: Basic Models Tests (Other)
timeout_in_minutes: 45
@@ -44,10 +41,10 @@ steps:
device: mi325_1
depends_on:
- image-build-amd
- label: Basic Models Test (Other CPU) # 5min
depends_on:
depends_on:
- image-build-cpu
timeout_in_minutes: 10
source_file_dependencies:
@@ -72,18 +69,3 @@ steps:
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- label: Transformers Backward Compatibility Models Test
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
commands:
- pip install transformers==4.57.5
- pytest -v -s tests/models/test_initialization.py
- pytest -v -s tests/models/test_transformers.py
- pytest -v -s tests/models/multimodal/processing/
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/offline_inference/basic/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
+7 -12
View File
@@ -1,9 +1,10 @@
group: Models - Language
depends_on:
depends_on:
- image-build
steps:
- label: Language Models Tests (Standard)
timeout_in_minutes: 25
torch_nightly: true
source_file_dependencies:
- vllm/
- tests/models/language
@@ -11,11 +12,10 @@ steps:
# Test standard language models, excluding a subset of slow tests
- pip freeze | grep -E 'torch'
- pytest -v -s models/language -m 'core_model and (not slow_test)'
mirror:
torch_nightly: {}
- label: Language Models Tests (Extra Standard) %N
timeout_in_minutes: 45
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
- tests/models/language/pooling/test_embedding.py
@@ -27,11 +27,10 @@ steps:
- pip freeze | grep -E 'torch'
- pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
parallelism: 2
mirror:
torch_nightly: {}
- label: Language Models Tests (Hybrid) %N
timeout_in_minutes: 75
torch_nightly: true
source_file_dependencies:
- vllm/
- tests/models/language/generation
@@ -39,12 +38,10 @@ steps:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
# Shard hybrid language model tests
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
parallelism: 2
mirror:
torch_nightly: {}
- label: Language Models Test (Extended Generation) # 80min
timeout_in_minutes: 110
@@ -56,7 +53,7 @@ steps:
# Install fast path packages for testing against transformers
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
mirror:
amd:
@@ -65,12 +62,11 @@ steps:
- image-build-amd
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
- label: Language Models Test (PPL)
timeout_in_minutes: 110
device: h200_18gb
optional: true
source_file_dependencies:
- vllm/
@@ -94,7 +90,6 @@ steps:
- label: Language Models Test (MTEB)
timeout_in_minutes: 110
device: h200_18gb
optional: true
source_file_dependencies:
- vllm/
+1 -7
View File
@@ -4,7 +4,6 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -20,7 +19,6 @@ steps:
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -28,7 +26,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
@@ -57,8 +54,7 @@ steps:
- tests/models/multimodal
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
@@ -81,7 +77,6 @@ steps:
- label: Multi-Modal Processor # 44min
timeout_in_minutes: 60
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -136,7 +131,6 @@ steps:
- label: Multi-Modal Models (Extended Pooling)
optional: true
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal/pooling
+1 -3
View File
@@ -49,7 +49,6 @@ steps:
- label: PyTorch Fullgraph
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/compile
@@ -61,9 +60,8 @@ steps:
# if this test fails, it means the nightly torch version is not compatible with some
# of the dependencies. Please check the error message and add the package to whitelist
# in /vllm/tools/pre_commit/generate_nightly_torch_test.py
device: h200_18gb
soft_fail: true
source_file_dependencies:
- requirements/test/nightly-torch.txt
- requirements/nightly_torch_test.txt
commands:
- bash standalone_tests/pytorch_nightly_dependency.sh
+3 -3
View File
@@ -1,5 +1,5 @@
group: Quantization
depends_on:
depends_on:
- image-build
steps:
- label: Quantization
@@ -9,14 +9,14 @@ steps:
- vllm/model_executor/layers/quantization
- tests/quantization
commands:
# temporary install here since we need nightly, will move to requirements/test/cuda.in
# temporary install here since we need nightly, will move to requirements/test.in
# after torchao 0.12 release, and pin a working version of torchao nightly here
# since torchao nightly is only compatible with torch nightly currently
# https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now
# we can only upgrade after this is resolved
# TODO(jerryzh168): resolve the above comment
- uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
- uv pip install --system torchao==0.14.1 --index-url https://download.pytorch.org/whl/cu129
- uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
-1
View File
@@ -7,7 +7,6 @@ steps:
# If this fails, it means the PR introduces a dependency that
# conflicts with Ray's dependency constraints.
# See https://github.com/vllm-project/vllm/issues/33599
device: h200_18gb
soft_fail: true
timeout_in_minutes: 10
source_file_dependencies:
-51
View File
@@ -4,18 +4,6 @@ depends_on:
steps:
- label: Spec Decode Eagle
timeout_in_minutes: 30
device: h200_18gb
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 Eagle Nightly B200
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
@@ -25,7 +13,6 @@ steps:
- label: Spec Decode Speculators + MTP
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
@@ -34,21 +21,8 @@ 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
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
@@ -58,34 +32,9 @@ steps:
- label: Spec Decode Draft Model
timeout_in_minutes: 30
device: h200_18gb
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: 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
optional: true
num_devices: 1
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- tests/v1/spec_decode/test_speculators_dflash.py
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
+10 -14
View File
@@ -3,7 +3,7 @@
# This lists cover the "core" components of vLLM that require careful review
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
@@ -44,9 +44,8 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/pooling_params.py @noooop @DarkLight1337
/vllm/tokenizers @DarkLight1337 @njhill
/vllm/renderers @DarkLight1337 @njhill
/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/reasoning @aarnphm @chaunceyjiang
/vllm/tool_parsers @aarnphm @chaunceyjiang
# vLLM V1
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -92,10 +91,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/tests/v1/kv_connector/nixl_integration @NickLucche
/tests/v1/kv_connector @ApostaC @orozery
/tests/v1/kv_offload @ApostaC @orozery
/tests/v1/determinism @yewentao256
/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/v1/determinism @yewentao256
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
@@ -124,16 +120,16 @@ mkdocs.yaml @hmellor
/tools/pre_commit @hmellor
# CPU
/vllm/v1/worker/cpu* @bigPYJ1151 @xuechendi
/vllm/v1/worker/cpu* @bigPYJ1151
/csrc/cpu @bigPYJ1151
/vllm/platforms/cpu.py @bigPYJ1151 @xuechendi
/vllm/platforms/cpu.py @bigPYJ1151
/cmake/cpu_extension.cmake @bigPYJ1151
/docker/Dockerfile.cpu @bigPYJ1151 @xuechendi
/docker/Dockerfile.cpu @bigPYJ1151
# Intel GPU
/vllm/v1/worker/xpu* @jikunshang @xuechendi
/vllm/platforms/xpu.py @jikunshang @xuechendi
/docker/Dockerfile.xpu @jikunshang @xuechendi
/vllm/v1/worker/xpu* @jikunshang
/vllm/platforms/xpu.py @jikunshang
/docker/Dockerfile.xpu @jikunshang
# Nemotron-specific files
/vllm/model_executor/models/*nemotron* @tomeras91
+1
View File
@@ -15,6 +15,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
</details>
**BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)
+10 -32
View File
@@ -18,7 +18,7 @@ pull_request_rules:
- name: comment-pre-commit-failure
description: Comment on PR when pre-commit check fails
conditions:
- check-failure=pre-commit
- status-failure=pre-commit
- -closed
- -draft
actions:
@@ -51,7 +51,7 @@ pull_request_rules:
- name: comment-dco-failure
description: Comment on PR when DCO check fails
conditions:
- check-failure=dco
- status-failure=dco
- -closed
- -draft
actions:
@@ -83,8 +83,8 @@ pull_request_rules:
- or:
- files~=^examples/.*deepseek.*\.py
- files~=^tests/.*deepseek.*\.py
- files~=^vllm/entrypoints/openai/tool_parsers/.*deepseek.*\.py
- files~=^vllm/model_executor/models/.*deepseek.*\.py
- files~=^vllm/tool_parsers/.*deepseek.*\.py
- files~=^vllm/reasoning/.*deepseek.*\.py
- files~=^vllm/transformers_utils/.*deepseek.*\.py
- title~=(?i)DeepSeek
@@ -110,10 +110,9 @@ pull_request_rules:
- or:
- files~=^examples/.*llama.*\.py
- files~=^tests/.*llama.*\.py
- files~=^vllm/entrypoints/openai/tool_parsers/llama.*\.py
- files~=^vllm/model_executor/models/.*llama.*\.py
- files~=^vllm/reasoning/.*llama.*\.py
- files~=^vllm/tool_parsers/.*llama.*\.py
- files~=^vllm/transformers_utils/.*llama.*\.py
- files~=^vllm/transformers_utils/configs/.*llama.*\.py
- title~=(?i)llama
actions:
label:
@@ -134,23 +133,6 @@ pull_request_rules:
add:
- multi-modality
- name: label-mistral
description: Automatically apply mistral label
conditions:
- label != stale
- or:
- files~=^examples/.*mistral.*\.py
- files~=^tests/.*mistral.*\.py
- files~=^vllm/model_executor/models/.*mistral.*\.py
- files~=^vllm/reasoning/.*mistral.*\.py
- files~=^vllm/tool_parsers/.*mistral.*\.py
- files~=^vllm/transformers_utils/.*mistral.*\.py
- title~=(?i)Mistral
actions:
label:
add:
- mistral
- name: label-new-model
description: Automatically apply new-model label
conditions:
@@ -185,9 +167,7 @@ pull_request_rules:
- files~=^examples/.*qwen.*\.py
- files~=^tests/.*qwen.*\.py
- files~=^vllm/model_executor/models/.*qwen.*\.py
- files~=^vllm/tool_parsers/.*qwen.*\.py
- files~=^vllm/reasoning/.*qwen.*\.py
- files~=^vllm/transformers_utils/.*qwen.*\.py
- title~=(?i)Qwen
actions:
label:
@@ -262,9 +242,8 @@ 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
- files=vllm/distributed/device_communicators/xpu_communicator.py
- files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -272,7 +251,6 @@ pull_request_rules:
- files=vllm/v1/worker/xpu_worker.py
- files=vllm/v1/worker/xpu_model_runner.py
- files=vllm/_xpu_ops.py
- files=vllm/kernels/xpu_ops.py
- files~=^vllm/lora/ops/xpu_ops
- files=vllm/lora/punica_wrapper/punica_xpu.py
- files=vllm/platforms/xpu.py
@@ -280,6 +258,7 @@ pull_request_rules:
- title~=(?i)XPU
- title~=(?i)Intel
- title~=(?i)BMG
- title~=(?i)Arc
actions:
label:
add:
@@ -399,18 +378,17 @@ pull_request_rules:
add:
- tool-calling
- name: auto-rebase to keep merge candidate within 1 day behind main
- name: auto-rebase if approved, ready, and 40 commits behind main
conditions:
- base = main
- label=ready
- "#approved-reviews-by >= 1"
- "#commits-behind >= 50"
- "#check-failure = 0"
- "#commits-behind >= 40"
- -closed
- -draft
- -conflict
actions:
update: {}
rebase: {}
- name: ping author on conflicts and add 'needs-rebase' label
conditions:
+4 -9
View File
@@ -320,25 +320,20 @@ jobs:
script: |
// Configuration: Map labels to GitHub users to CC
// You can add multiple users per label, and multiple label configurations
// {users} will be replaced with @mentions
const ccConfig = {
rocm: {
users: ['hongxiayang', 'tjtanaa', 'vllmellm'],
message: 'CC {users} for ROCm-related issue',
},
mistral: {
users: ['patrickvonplaten', 'juliendenize', 'andylolu2'],
message: 'CC {users} for Mistral-related issue',
users: ['hongxiayang', 'tjtanaa', 'vllmellm'], // Add more users as needed: ['user1', 'user2', 'user3']
message: 'CC {users} for ROCm-related issue' // {users} will be replaced with @mentions
},
// Add more label -> user mappings here
// Example:
// cuda: {
// users: ['user1', 'user2'],
// message: 'CC {users} for CUDA-related issue',
// message: 'CC {users} for CUDA-related issue'
// },
// performance: {
// users: ['perfexpert'],
// message: 'CC {users} for performance issue',
// message: 'CC {users} for performance issue'
// },
};
+1 -2
View File
@@ -32,7 +32,7 @@ jobs:
- name: Install dependencies and build vLLM
run: |
uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match
uv pip install -r requirements/cpu-build.txt --index-strategy unsafe-best-match
uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
uv pip install -e . --no-build-isolation
env:
@@ -45,7 +45,6 @@ jobs:
- name: Smoke test vllm serve
run: |
# Start server in background
VLLM_CPU_KVCACHE_SPACE=1 \
vllm serve Qwen/Qwen3-0.6B \
--max-model-len=2K \
--load-format=dummy \
+5 -5
View File
@@ -62,14 +62,14 @@ jobs:
const prAuthor = context.payload.pull_request.user.login;
const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
per_page: 1,
});
const mergedPRCount = searchResults.total_count;
console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
const authorPRCount = searchResults.total_count;
console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
if (mergedPRCount === 0) {
if (authorPRCount === 1) {
console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
await github.rest.issues.createComment({
owner,
@@ -98,5 +98,5 @@ jobs:
].join('\n'),
});
} else {
console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
}
+2 -12
View File
@@ -2,7 +2,6 @@ name: pre-commit
on:
pull_request:
types: [opened, synchronize, reopened, labeled]
push:
branches: [main]
@@ -16,11 +15,7 @@ permissions:
jobs:
pre-run-check:
if: >-
github.event_name == 'pull_request' &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified')
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
@@ -49,12 +44,7 @@ jobs:
pre-commit:
needs: pre-run-check
if: >-
always() &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified') &&
(needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+2 -2
View File
@@ -9,12 +9,12 @@ PATH=${cuda_home}/bin:$PATH
LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
# Install requirements
$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
$python_executable -m pip install -r requirements/build.txt -r requirements/cuda.txt
# 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
-4
View File
@@ -12,9 +12,6 @@ vllm/third_party/triton_kernels/*
# FlashMLA interface copied from source
vllm/third_party/flashmla/flash_mla_interface.py
# DeepGEMM vendored package built from source
vllm/third_party/deep_gemm/
# triton jit
.triton
@@ -29,7 +26,6 @@ __pycache__/
# Distribution / packaging
.Python
build/
!requirements/build/
cmake-build-*/
CMakeUserPresets.json
develop-eggs/
+10 -65
View File
@@ -39,24 +39,15 @@ repos:
rev: 0.11.1
hooks:
- id: pip-compile
args: [
requirements/test/cuda.in,
-c, requirements/cuda.txt,
-o, requirements/test/cuda.txt,
--index-strategy, unsafe-best-match,
--torch-backend, cu130,
--python-platform, x86_64-manylinux_2_28,
--python-version, "3.12",
]
files: ^requirements/(common|cuda|test/cuda)\.(in|txt)$
args: [requirements/test.in, -c, requirements/common.txt, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
files: ^requirements/test\.(in|txt)$
- id: pip-compile
alias: pip-compile-rocm
name: pip-compile-rocm
args: [
requirements/test/rocm.in,
-c, requirements/rocm.txt,
-o, requirements/test/rocm.txt,
requirements/rocm-test.in, -o, requirements/rocm-test.txt,
--index-strategy, unsafe-best-match,
-c, requirements/rocm.txt,
--python-platform, x86_64-manylinux_2_28,
--python-version, "3.12",
# Exclude torch and CUDA/NVIDIA packages
@@ -68,76 +59,30 @@ repos:
--no-emit-package, cuda-pathfinder,
--no-emit-package, cuda-toolkit,
--no-emit-package, cupy-cuda12x,
# nvidia packages (unsuffixed / unified naming)
--no-emit-package, nvidia-cublas,
--no-emit-package, nvidia-cuda-cupti,
--no-emit-package, nvidia-cuda-nvrtc,
--no-emit-package, nvidia-cuda-runtime,
--no-emit-package, nvidia-cudnn,
--no-emit-package, nvidia-cudnn-cu13,
--no-emit-package, nvidia-cufft,
--no-emit-package, nvidia-cufile,
--no-emit-package, nvidia-curand,
--no-emit-package, nvidia-cusolver,
--no-emit-package, nvidia-cusparse,
--no-emit-package, nvidia-cusparselt,
--no-emit-package, nvidia-nccl,
--no-emit-package, nvidia-nvjitlink,
--no-emit-package, nvidia-nvshmem,
--no-emit-package, nvidia-nvtx,
# nvidia cu12 packages
--no-emit-package, nvidia-cublas-cu12,
--no-emit-package, nvidia-cuda-cupti-cu12,
--no-emit-package, nvidia-cuda-nvrtc-cu12,
--no-emit-package, nvidia-cuda-runtime-cu12,
--no-emit-package, nvidia-cudnn-cu12,
--no-emit-package, nvidia-cufft-cu12,
--no-emit-package, nvidia-cufile-cu12,
--no-emit-package, nvidia-curand-cu12,
--no-emit-package, nvidia-cusolver-cu12,
--no-emit-package, nvidia-cusparse-cu12,
--no-emit-package, nvidia-cusparselt-cu12,
--no-emit-package, nvidia-nccl-cu12,
--no-emit-package, nvidia-nvjitlink-cu12,
--no-emit-package, nvidia-nvshmem-cu12,
--no-emit-package, nvidia-nvtx-cu12,
# nvidia cu13 packages
--no-emit-package, nvidia-cublas-cu13,
--no-emit-package, nvidia-cuda-cupti-cu13,
--no-emit-package, nvidia-cuda-nvrtc-cu13,
--no-emit-package, nvidia-cuda-runtime-cu13,
--no-emit-package, nvidia-cudnn-cu13,
--no-emit-package, nvidia-cufft-cu13,
--no-emit-package, nvidia-cufile-cu13,
--no-emit-package, nvidia-curand-cu13,
--no-emit-package, nvidia-cusolver-cu13,
--no-emit-package, nvidia-cusparse-cu13,
--no-emit-package, nvidia-cusparselt-cu13,
--no-emit-package, nvidia-nccl-cu13,
--no-emit-package, nvidia-nvjitlink-cu13,
--no-emit-package, nvidia-nvjitlink,
--no-emit-package, nvidia-nvshmem-cu13,
--no-emit-package, nvidia-nvtx-cu13,
--no-emit-package, nvidia-nvtx,
]
files: ^requirements/(common|rocm|test/rocm)\.(in|txt)$
- id: pip-compile
alias: pip-compile-xpu
name: pip-compile-xpu
args: [
requirements/test/xpu.in,
-c, requirements/xpu.txt,
-o, requirements/test/xpu.txt,
--index-strategy, unsafe-best-match,
--torch-backend, xpu,
--python-platform, x86_64-manylinux_2_39,
--python-version, "3.12",
]
files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$
files: ^requirements/rocm-test\.(in|txt)$
- repo: local
hooks:
- id: format-torch-nightly-test
name: reformat test/nightly-torch.txt to be in sync with test/cuda.in
name: reformat nightly_torch_test.txt to be in sync with test.in
language: python
entry: python tools/pre_commit/generate_nightly_torch_test.py
files: ^requirements/test/cuda\.(in|txt)$
files: ^requirements/test\.(in|txt)$
- id: mypy-local
name: Run mypy locally for lowest supported Python version
entry: python tools/pre_commit/mypy.py 0 "3.10"
+3 -3
View File
@@ -72,11 +72,11 @@ uv pip install -e . --torch-backend=auto
```bash
# Install test dependencies.
# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the
# requirements/test.txt is pinned to x86_64; on other platforms, use the
# unpinned source file instead:
uv pip install -r requirements/test/cuda.in # resolves for current platform
uv pip install -r requirements/test.in # resolves for current platform
# Or on x86_64:
uv pip install -r requirements/test/cuda.txt
uv pip install -r requirements/test.txt
# Run a specific test file (use .venv/bin/python directly;
# `source activate` does not persist in non-interactive shells):
+12 -34
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.
@@ -56,8 +56,8 @@ endif()
# requirements.txt files and should be kept consistent. The ROCm torch
# versions are derived from docker/Dockerfile.rocm
#
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
set(TORCH_SUPPORTED_VERSION_CUDA "2.10.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.10.0")
#
# Try to find python package with an executable that exactly matches
@@ -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()
#
@@ -228,8 +225,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
# Certain HIP functions are marked as [[nodiscard]], yet vllm ignores the result which generates
# a lot of warnings that always mask real issues. Suppressing until this is properly addressed.
#
set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result -Wno-unused-value")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result -Wno-unused-value")
set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")
endif()
#
@@ -302,7 +299,6 @@ set(VLLM_EXT_SRC
"csrc/quantization/w8a8/int8/scaled_quant.cu"
"csrc/quantization/w8a8/fp8/common.cu"
"csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
"csrc/quantization/gguf/gguf_kernel.cu"
"csrc/quantization/activation_kernels.cu"
"csrc/cuda_utils_kernels.cu"
@@ -310,8 +306,6 @@ 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")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
@@ -346,7 +340,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_EXT_SRC
"csrc/quantization/awq/gemm_kernels.cu"
"csrc/cutlass_extensions/common.cpp")
"csrc/cutlass_extensions/common.cpp"
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu")
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
@@ -926,14 +921,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -955,19 +942,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -1243,7 +1222,6 @@ endif()
# For CUDA we also build and ship some external projects.
if (VLLM_GPU_LANG STREQUAL "CUDA")
include(cmake/external_projects/deepgemm.cmake)
include(cmake/external_projects/flashmla.cmake)
include(cmake/external_projects/qutlass.cmake)
+20 -27
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.
---
@@ -23,54 +23,47 @@ For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
vLLM is a fast and easy-to-use library for LLM inference and serving.
Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors.
Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has evolved into a community-driven project with contributions from both academia and industry.
vLLM is fast with:
- State-of-the-art serving throughput
- Efficient management of attention key and value memory with [**PagedAttention**](https://blog.vllm.ai/2023/06/20/vllm.html)
- Continuous batching of incoming requests, chunked prefill, prefix caching
- Fast and flexible model execution with piecewise and full CUDA/HIP graphs
- Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and [more](https://docs.vllm.ai/en/latest/features/quantization/index.html)
- Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton
- Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL
- Speculative decoding including n-gram, suffix, EAGLE, DFlash
- Automatic kernel generation and graph-level transformations using torch.compile
- Disaggregated prefill, decode, and encode
- Continuous batching of incoming requests
- Fast model execution with CUDA/HIP graph
- Quantizations: [GPTQ](https://arxiv.org/abs/2210.17323), [AWQ](https://arxiv.org/abs/2306.00978), [AutoRound](https://arxiv.org/abs/2309.05516), INT4, INT8, and FP8
- Optimized CUDA kernels, including integration with FlashAttention and FlashInfer
- Speculative decoding
- Chunked prefill
vLLM is flexible and easy to use with:
- Seamless integration with popular Hugging Face models
- High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more
- Tensor, pipeline, data, expert, and context parallelism for distributed inference
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- Generation of structured outputs using xgrammar or guidance
- Tool calling and reasoning parsers
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
- 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.
- OpenAI-compatible API server
- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, Arm CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend.
- Prefix caching support
- Multi-LoRA support
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
vLLM seamlessly supports most popular open-source models on HuggingFace, including:
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
- Hybrid attention and state-space models (e.g., Mamba, Qwen3.5)
- Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral)
- Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT)
- Reward and classification models (e.g., Qwen-Math)
- Transformer-like LLMs (e.g., Llama)
- Mixture-of-Expert LLMs (e.g., Mixtral, Deepseek-V2 and V3)
- Embedding Models (e.g., E5-Mistral)
- Multi-modal LLMs (e.g., LLaVA)
Find the full list of supported models [here](https://docs.vllm.ai/en/latest/models/supported_models.html).
## Getting Started
Install vLLM with [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip`:
Install vLLM with `pip` or [from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source):
```bash
uv pip install vllm
pip install vllm
```
Or [build from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for development.
Visit our [documentation](https://docs.vllm.ai/en/latest/) to learn more.
- [Installation](https://docs.vllm.ai/en/latest/getting_started/installation.html)
@@ -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
+7 -12
View File
@@ -9,12 +9,11 @@ os.environ["VLLM_USE_DEEP_GEMM"] = "0"
import torch
from vllm.benchmarks.lib.utils import default_vllm_config
from vllm.model_executor.kernels.linear import (
init_fp8_linear_kernel,
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
W8A8BlockFp8LinearOp,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
create_fp8_quant_key,
)
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
CUTLASS_BLOCK_FP8_SUPPORTED,
@@ -71,15 +70,11 @@ def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass):
weight_group_shape = GroupShape(block_n, block_k)
act_quant_group_shape = GroupShape(1, block_k) # Per-token, per-group quantization
linear_op = init_fp8_linear_kernel(
weight_quant_key=create_fp8_quant_key(
static=True, group_shape=weight_group_shape
),
activation_quant_key=create_fp8_quant_key(
static=False, group_shape=act_quant_group_shape
),
out_dtype=torch.get_default_dtype(),
module_name="build_w8a8_block_fp8_runner",
linear_op = W8A8BlockFp8LinearOp(
weight_group_shape=weight_group_shape,
act_quant_group_shape=act_quant_group_shape,
cutlass_block_fp8_supported=use_cutlass,
use_aiter_and_is_supported=False,
)
def run():
@@ -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,
@@ -9,7 +9,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
moe_align_block_size,
)
from vllm.triton_utils import triton
from vllm.utils.torch_utils import set_random_seed
def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
@@ -45,7 +44,7 @@ configs = list(
def benchmark(num_tokens, num_experts, topk, ep_size, provider):
"""Benchmark function for Triton."""
block_size = 256
set_random_seed(0)
torch.cuda.manual_seed_all(0)
topk_ids = get_topk_ids(num_tokens, num_experts, topk)
e_map = None
@@ -20,7 +20,7 @@ import matplotlib.pyplot as plt
import numpy as np
import torch
from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
persistent_masked_m_silu_mul_quant,
)
from vllm.triton_utils import tl, triton
@@ -16,7 +16,6 @@ from vllm.utils.deep_gemm import (
fp8_gemm_nt,
per_block_cast_to_fp8,
)
from vllm.utils.torch_utils import set_random_seed
def benchmark_shape(
@@ -236,7 +235,9 @@ def run_benchmarks(verbose: bool = False):
torch.backends.cudnn.allow_tf32 = True
# Set seeds for reproducibility
set_random_seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
# Define benchmark shapes (m, n, k)
shapes = [
(8, 4096, 7168),
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
],
}
@@ -1439,12 +1439,6 @@ async def main() -> None:
action="store_true",
help="Export summary to Excel file (optional)",
)
parser.add_argument(
"--stats-json-output",
type=str,
default=None,
help="Export per-request stats (ttft_ms, tpot_ms, etc.) to a JSON file",
)
parser.add_argument(
"-v",
"--verbose",
@@ -1657,19 +1651,6 @@ async def main() -> None:
warmup_runtime_sec=warmup_runtime_sec,
)
if args.stats_json_output is not None:
# Export per-request metrics as a JSON array for downstream analysis.
stats_data = [s._asdict() for s in client_metrics]
logger.info(
f"{Color.GREEN}Writing per-request stats JSON: "
f"{args.stats_json_output}{Color.RESET}"
)
os.makedirs(
os.path.dirname(os.path.abspath(args.stats_json_output)), exist_ok=True
)
with open(args.stats_json_output, "w") as f:
json.dump(stats_data, f, indent=2)
if args.output_file is not None:
# Write a JSON file with the updated conversations
# The "assistant" content will contain the answers from the tested LLM
+22 -60
View File
@@ -30,21 +30,6 @@ else()
list(APPEND CXX_COMPILE_FLAGS
"-fopenmp"
"-DVLLM_CPU_EXTENSION")
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
endif()
if (NOT MACOSX_FOUND)
@@ -161,49 +146,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()
@@ -223,6 +175,20 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
if(NOT NPROC)
set(NPROC 4)
endif()
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
# Fetch and populate ACL
if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
@@ -383,7 +349,6 @@ endif()
set(VLLM_EXT_SRC
"csrc/cpu/activation.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/layernorm.cpp"
"csrc/cpu/mla_decode.cpp"
"csrc/cpu/pos_encoding.cpp"
@@ -394,7 +359,6 @@ set(VLLM_EXT_SRC
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
"csrc/cpu/activation_lut_bf16.cpp"
${VLLM_EXT_SRC})
endif()
@@ -419,7 +383,6 @@ if (ENABLE_X86_ISA)
"csrc/cpu/cpu_wna16.cpp"
"csrc/cpu/cpu_fused_moe.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/torch_bindings.cpp"
@@ -432,7 +395,6 @@ if (ENABLE_X86_ISA)
set(VLLM_EXT_SRC_AVX2
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
-151
View File
@@ -1,151 +0,0 @@
include(FetchContent)
# If DEEPGEMM_SRC_DIR is set, DeepGEMM is built from that directory
# instead of downloading.
# It can be set as an environment variable or passed as a cmake argument.
# The environment variable takes precedence.
if (DEFINED ENV{DEEPGEMM_SRC_DIR})
set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
endif()
if(DEEPGEMM_SRC_DIR)
FetchContent_Declare(
deepgemm
SOURCE_DIR ${DEEPGEMM_SRC_DIR}
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
else()
# This ref should be kept in sync with tools/install_deepgemm.sh
FetchContent_Declare(
deepgemm
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
)
endif()
# Use FetchContent_Populate (not MakeAvailable) to avoid processing
# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
FetchContent_GetProperties(deepgemm)
if(NOT deepgemm_POPULATED)
FetchContent_Populate(deepgemm)
endif()
message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
set(DEEPGEMM_SUPPORT_ARCHS)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
endif()
cuda_archs_loose_intersection(DEEPGEMM_ARCHS
"${DEEPGEMM_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
if(DEEPGEMM_ARCHS)
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
find_package(CUDAToolkit REQUIRED)
#
# Build the _C pybind11 extension from DeepGEMM's C++ source.
# This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
#
Python_add_library(_deep_gemm_C MODULE WITH_SOABI
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
# The pybind11 module name must be _C to match DeepGEMM's Python imports.
set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
target_compile_definitions(_deep_gemm_C PRIVATE
"-DTORCH_EXTENSION_NAME=_C")
target_include_directories(_deep_gemm_C PRIVATE
"${deepgemm_SOURCE_DIR}/csrc"
"${deepgemm_SOURCE_DIR}/deep_gemm/include"
"${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
"${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
"${deepgemm_SOURCE_DIR}/third-party/fmt/include")
target_compile_options(_deep_gemm_C PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
$<$<COMPILE_LANGUAGE:CXX>:-O3>
$<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
$<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
# torch_python is required because DeepGEMM uses pybind11 type casters
# for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
# use torch::Library custom ops.
find_library(TORCH_PYTHON_LIBRARY torch_python
PATHS "${TORCH_INSTALL_PREFIX}/lib"
REQUIRED)
target_link_libraries(_deep_gemm_C PRIVATE
torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
CUDA::cudart CUDA::nvrtc)
# Install the shared library into the vendored package directory
install(TARGETS _deep_gemm_C
LIBRARY DESTINATION vllm/third_party/deep_gemm
COMPONENT _deep_gemm_C)
#
# Vendor DeepGEMM Python package files
#
install(FILES
"${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py"
DESTINATION vllm/third_party/deep_gemm
COMPONENT _deep_gemm_C)
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/utils/"
DESTINATION vllm/third_party/deep_gemm/utils
COMPONENT _deep_gemm_C
FILES_MATCHING PATTERN "*.py")
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/testing/"
DESTINATION vllm/third_party/deep_gemm/testing
COMPONENT _deep_gemm_C
FILES_MATCHING PATTERN "*.py")
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/legacy/"
DESTINATION vllm/third_party/deep_gemm/legacy
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")
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
DESTINATION vllm/third_party/deep_gemm
RENAME envs.py
COMPONENT _deep_gemm_C)
#
# Install include files needed for JIT compilation at runtime.
# The JIT compiler finds these relative to the package directory.
#
# DeepGEMM's own CUDA headers
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/include/"
DESTINATION vllm/third_party/deep_gemm/include
COMPONENT _deep_gemm_C)
# CUTLASS and CuTe headers (vendored for JIT, separate from vLLM's CUTLASS)
install(DIRECTORY "${deepgemm_SOURCE_DIR}/third-party/cutlass/include/"
DESTINATION vllm/third_party/deep_gemm/include
COMPONENT _deep_gemm_C)
else()
message(STATUS "DeepGEMM will not compile: "
"unsupported CUDA architecture ${CUDA_ARCHS}")
# Create empty target so setup.py doesn't fail on unsupported systems
add_custom_target(_deep_gemm_C)
endif()
+16 -28
View File
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
GIT_TAG c0ec424fd8a546d0cbbf4bf050bbcfe837c55afb
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -87,30 +87,18 @@ endforeach()
#
add_custom_target(_vllm_fa4_cutedsl_C)
# Install flash_attn/cute directory (needed for FA4).
# When using a local source dir (VLLM_FLASH_ATTN_SRC_DIR), create a symlink
# so edits to cute-dsl Python files take effect immediately without rebuilding.
# Otherwise, copy files and transform flash_attn.cute imports to
# vllm.vllm_flash_attn.cute to match our package structure.
if(VLLM_FLASH_ATTN_SRC_DIR)
install(CODE "
set(LINK_TARGET \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\")
set(LINK_NAME \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute\")
file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")
file(REMOVE_RECURSE \"\${LINK_NAME}\")
file(CREATE_LINK \"\${LINK_TARGET}\" \"\${LINK_NAME}\" SYMBOLIC)
" COMPONENT _vllm_fa4_cutedsl_C)
else()
install(CODE "
file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
foreach(SRC_FILE \${CUTE_PY_FILES})
file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
file(MAKE_DIRECTORY \${DST_DIR})
file(READ \${SRC_FILE} FILE_CONTENTS)
string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
endforeach()
" COMPONENT _vllm_fa4_cutedsl_C)
endif()
# Copy flash_attn/cute directory (needed for FA4) and transform imports
# The cute directory uses flash_attn.cute imports internally, which we replace
# with vllm.vllm_flash_attn.cute to match our package structure.
install(CODE "
file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")
foreach(SRC_FILE \${CUTE_PY_FILES})
file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})
set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")
get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)
file(MAKE_DIRECTORY \${DST_DIR})
file(READ \${SRC_FILE} FILE_CONTENTS)
string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")
file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")
endforeach()
" COMPONENT _vllm_fa4_cutedsl_C)
-100
View File
@@ -1,100 +0,0 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
*
* 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.
*/
#pragma once
namespace vllm {
namespace cuda_async {
__device__ __forceinline__ void cp_async_shared_global_16_cg(
void* smem_ptr, const void* glob_ptr) {
#if defined(USE_ROCM)
*reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
:
: "r"(smem), "l"(glob_ptr));
#elif defined(__CUDA_ARCH__)
*reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
#else
(void)smem_ptr;
(void)glob_ptr;
#endif
}
__device__ __forceinline__ void cp_async_shared_global_ca(void* smem_ptr,
const void* glob_ptr,
int size_bytes) {
#if defined(USE_ROCM)
if (size_bytes == 4) {
*reinterpret_cast<uint32_t*>(smem_ptr) =
*reinterpret_cast<const uint32_t*>(glob_ptr);
} else if (size_bytes == 8) {
*reinterpret_cast<uint64_t*>(smem_ptr) =
*reinterpret_cast<const uint64_t*>(glob_ptr);
} else {
*reinterpret_cast<int4*>(smem_ptr) =
*reinterpret_cast<const int4*>(glob_ptr);
}
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
if (size_bytes == 4) {
asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n"
:
: "r"(smem), "l"(glob_ptr));
} else if (size_bytes == 8) {
asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n"
:
: "r"(smem), "l"(glob_ptr));
} else {
asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n"
:
: "r"(smem), "l"(glob_ptr));
}
#elif defined(__CUDA_ARCH__)
if (size_bytes == 4) {
*reinterpret_cast<uint32_t*>(smem_ptr) =
*reinterpret_cast<const uint32_t*>(glob_ptr);
} else if (size_bytes == 8) {
*reinterpret_cast<uint64_t*>(smem_ptr) =
*reinterpret_cast<const uint64_t*>(glob_ptr);
} else {
*reinterpret_cast<int4*>(smem_ptr) =
*reinterpret_cast<const int4*>(glob_ptr);
}
#else
(void)smem_ptr;
(void)glob_ptr;
(void)size_bytes;
#endif
}
__device__ __forceinline__ void cp_async_commit_group() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
asm volatile("cp.async.commit_group;\n" ::);
#endif
}
template <int n>
__device__ __forceinline__ void cp_async_wait_group() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
asm volatile("cp.async.wait_group %0;\n" : : "n"(n));
#endif
}
} // namespace cuda_async
} // namespace vllm
-16
View File
@@ -17,22 +17,6 @@ enum class Fp8KVCacheDataType {
kFp8E5M2 = 2,
};
inline Fp8KVCacheDataType get_fp8_kv_cache_data_type(
const std::string& dtype_str) {
// dtype_str refers to CacheDType at vllm.config.cache.CacheDType
if (dtype_str == "auto" || dtype_str == "float16" ||
dtype_str == "bfloat16") {
// unquantized kv cache
return Fp8KVCacheDataType::kAuto;
} else if (dtype_str == "fp8" || dtype_str == "fp8_ds_mla" ||
dtype_str == "fp8_e4m3") {
return Fp8KVCacheDataType::kFp8E4M3;
} else if (dtype_str == "fp8_e5m2") {
return Fp8KVCacheDataType::kFp8E5M2;
}
TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str);
}
// fp8 vector types for quantization of kv cache
template <>
struct Vec<uint8_t, 1> {
+34 -85
View File
@@ -91,9 +91,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
if (n == 0) return;
int64_t* src_data = src_ptrs.mutable_data_ptr<int64_t>();
int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
const int64_t* src_data = src_ptrs.data_ptr<int64_t>();
const int64_t* dst_data = dst_ptrs.data_ptr<int64_t>();
const int64_t* size_data = sizes.data_ptr<int64_t>();
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
@@ -104,49 +104,28 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
static_assert(sizeof(size_t) == sizeof(int64_t));
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
// Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
// binaries compiled with CUDA 12.8+ still work on older drivers, and
// we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
// The function pointer is cached after the first call.
using BatchFn =
CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t,
CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream);
static BatchFn batch_fn = []() -> BatchFn {
CUdriverProcAddressQueryResult sym_status;
void* fn_ptr = nullptr;
CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080,
CU_GET_PROC_ADDRESS_DEFAULT, &sym_status);
if (res != CUDA_SUCCESS || fn_ptr == nullptr) {
return nullptr;
}
return reinterpret_cast<BatchFn>(fn_ptr);
}();
if (batch_fn != nullptr) {
CUmemcpyAttributes attr = {};
attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
size_t attrs_idx = 0;
size_t fail_idx = 0;
CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
reinterpret_cast<CUdeviceptr*>(src_data),
reinterpret_cast<size_t*>(size_data),
static_cast<size_t>(n), &attr, &attrs_idx, 1,
&fail_idx, static_cast<CUstream>(stream));
TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
fail_idx, " with error ", result);
} else
#endif
{
// Fallback for CUDA < 12.8, older drivers, and ROCm:
// individual async copies.
// cudaMemcpyDefault lets the driver infer direction from pointer types.
for (int64_t i = 0; i < n; i++) {
cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
reinterpret_cast<void*>(src_data[i]),
static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
stream);
}
CUmemcpyAttributes attr = {};
attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
size_t attrs_idx = 0;
size_t fail_idx = 0;
CUresult result = cuMemcpyBatchAsync(
reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx,
static_cast<CUstream>(stream));
TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
fail_idx, " with error ", result);
#else
// Fallback for CUDA < 12.8 and ROCm: individual async copies.
// cudaMemcpyDefault lets the driver infer direction from pointer types.
for (int64_t i = 0; i < n; i++) {
cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
reinterpret_cast<void*>(src_data[i]),
static_cast<size_t>(size_data[i]), cudaMemcpyDefault,
stream);
}
#endif
}
namespace vllm {
@@ -599,11 +578,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 +590,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;
@@ -731,28 +703,6 @@ void reshape_and_cache_flash(
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (kv_cache_dtype == "nvfp4") {
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
// NVFP4 dispatch is compiled separately for SM100+.
extern void reshape_and_cache_nvfp4_dispatch(
torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
torch::Tensor & value_cache, torch::Tensor & slot_mapping,
torch::Tensor & k_scale, torch::Tensor & v_scale);
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
slot_mapping, k_scale, v_scale);
return;
#else
TORCH_CHECK(false,
"NVFP4 KV cache requires SM100+ (Blackwell). "
"Please rebuild vllm with a Blackwell-compatible CUDA target.");
#endif
}
// Original FP8/auto path.
int block_size = key_cache.size(1);
int64_t key_stride = key.stride(0);
@@ -770,6 +720,8 @@ void reshape_and_cache_flash(
dim3 grid(num_tokens);
dim3 block(std::min(num_heads * head_size, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
CALL_RESHAPE_AND_CACHE_FLASH);
@@ -1497,9 +1449,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 +1460,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),
-71
View File
@@ -1,71 +0,0 @@
#include "cpu_types.hpp"
#include <array>
#include <cstdint>
#include <mutex>
#include <string>
#include <ATen/ops/empty.h>
#include <ATen/ops/gelu.h>
#include <c10/util/BFloat16.h>
constexpr uint32_t ActivationLutSize = 1u << 16;
at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
void maybe_init_activation_lut_bf16(
uint16_t* lut, std::once_flag& once,
at::Tensor (*activation)(const at::Tensor&)) {
std::call_once(once, [&]() {
auto lut_input =
at::empty({static_cast<int64_t>(ActivationLutSize)},
at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
auto* lut_input_ptr = lut_input.data_ptr<float>();
#pragma omp parallel for
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
}
auto lut_output = activation(lut_input);
const auto* lut_output_ptr = lut_output.data_ptr<float>();
#pragma omp parallel for
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
}
});
}
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const uint16_t* lut, const char* op_name) {
TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
": input must be bfloat16");
TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
": out must be bfloat16");
TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
const auto* src =
reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
const int64_t n = input.numel();
CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
#pragma omp parallel for
for (int64_t i = 0; i < n; ++i) {
dst[i] = lut[src[i]];
}
CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
}
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const std::string& activation) {
if (activation == "gelu") {
static std::array<uint16_t, ActivationLutSize> lut{};
static std::once_flag once;
maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
activation_lut_bf16(out, input, lut.data(), "gelu_lut");
return;
}
TORCH_CHECK(false, "Unsupported activation: ", activation);
}
+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 -31
View File
@@ -147,9 +147,6 @@ struct AttentionMetadata {
case ISA::NEON:
ss << "NEON, ";
break;
case ISA::VXE:
ss << "VXE, ";
break;
}
ss << "workitem_group_num: " << workitem_group_num
<< ", reduction_item_num: " << reduction_item_num
@@ -1152,11 +1149,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 +1198,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 +1255,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 +1269,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
+1 -1
View File
@@ -53,7 +53,7 @@ class TileGemm82 {
const int64_t ldb, const int64_t ldc,
const int32_t block_size, const int32_t dynamic_k_size,
const bool accum_c) {
static_assert(0 < M && M <= 8);
static_assert(0 < M <= 8);
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
kv_cache_t* __restrict__ curr_b_0 = b_tile;
+1 -1
View File
@@ -68,7 +68,7 @@ class TileGemm161 {
const int64_t ldb, const int64_t ldc,
const int32_t block_size, const int32_t dynamic_k_size,
const bool accum_c) {
static_assert(0 < M && M <= 16);
static_assert(0 < M <= 16);
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
kv_cache_t* __restrict__ curr_b_0 = b_tile;
+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 -1
View File
@@ -39,7 +39,7 @@ class TileGemm82 {
template <int32_t M>
static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
static_assert(0 < M && M <= 8);
static_assert(0 < M <= 8);
using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
scalar_t* __restrict__ curr_b_0 = b_ptr;
-409
View File
@@ -1,409 +0,0 @@
#include "cpu_types.hpp"
#include <algorithm>
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& valid_sampled_tokens_count,
const torch::Tensor& query_start_loc_gpu,
torch::Tensor& token_indices_to_sample,
torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) {
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
const int64_t* valid_count_ptr =
valid_sampled_tokens_count.data_ptr<int64_t>();
const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr<int32_t>();
int32_t* indices_out_ptr = token_indices_to_sample.data_ptr<int32_t>();
int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr<int64_t>();
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx;
int64_t num_valid_tokens = valid_count_ptr[req_idx];
int64_t num_rejected = 0;
if (num_draft_tokens > 0) {
num_rejected = num_draft_tokens + 1 - num_valid_tokens;
}
int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1;
int32_t index_to_sample = q_last_tok_idx - num_rejected;
indices_out_ptr[req_idx] = index_to_sample;
rejected_out_ptr[req_idx] = num_rejected;
}
}
void eagle_prepare_next_token_padded_kernel_impl(
const torch::Tensor& sampled_token_ids,
const torch::Tensor& discard_request_mask,
const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) {
const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr<int64_t>();
const bool* discard_mask_ptr = discard_request_mask.data_ptr<bool>();
const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr<int64_t>();
int64_t* next_ids_out_ptr = next_token_ids.data_ptr<int64_t>();
int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr<int64_t>();
const int64_t stride = sampled_token_ids.stride(0);
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride;
int64_t valid_count = 0;
int64_t last_valid_token = -1;
for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) {
int64_t token = row_ptr[pos];
if (token != -1 && token < vocab_size) {
valid_count++;
last_valid_token = token;
}
}
bool discard = discard_mask_ptr[req_idx];
if (discard) {
next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx];
valid_count_out_ptr[req_idx] = 0;
} else {
next_ids_out_ptr[req_idx] =
(valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx];
valid_count_out_ptr[req_idx] = valid_count;
}
}
}
void eagle_step_slot_mapping_metadata_kernel_impl(
const torch::Tensor& positions, const torch::Tensor& block_table,
torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
torch::Tensor& out_slot_mapping, const int64_t block_size,
const int64_t max_model_len, const int64_t PAD_ID) {
const int64_t batch_size = positions.size(0);
const int64_t input_batch_size = out_slot_mapping.size(0);
const int64_t* pos_ptr = positions.data_ptr<int64_t>();
const int32_t* bt_ptr = block_table.data_ptr<int32_t>();
int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>();
int64_t* out_clamped_ptr = out_clamped_positions.data_ptr<int64_t>();
int64_t* out_slot_ptr = out_slot_mapping.data_ptr<int64_t>();
const int64_t bt_stride = block_table.stride(0);
const int64_t n_blocks_per_req = block_table.size(1);
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) {
if (req_idx >= batch_size) {
out_slot_ptr[req_idx] = PAD_ID;
continue;
}
int64_t position = pos_ptr[req_idx];
int64_t new_position = position + 1;
bool exceeds_max = new_position >= max_model_len;
int64_t clamped_position = exceeds_max ? 0 : new_position;
out_clamped_ptr[req_idx] = clamped_position;
int64_t block_number = clamped_position / block_size;
block_number = std::min(block_number, n_blocks_per_req - 1);
int32_t block_id = bt_ptr[req_idx * bt_stride + block_number];
int64_t slot_id = block_id * block_size + (clamped_position % block_size);
out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id;
int32_t seq_len = seq_lens_ptr[req_idx];
int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1);
new_seq_len = std::min(new_seq_len, static_cast<int32_t>(max_model_len));
seq_lens_ptr[req_idx] = new_seq_len;
}
}
void copy_and_expand_eagle_inputs_kernel_impl(
const torch::Tensor& target_token_ids,
const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
torch::Tensor& out_input_ids, torch::Tensor& out_positions,
torch::Tensor& out_is_rejected_token_mask,
torch::Tensor& out_is_masked_token_mask,
torch::Tensor& out_new_token_indices,
torch::Tensor& out_hidden_state_mapping,
const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
const int64_t total_input_tokens,
const int64_t num_padding_slots_per_request, const bool shift_input_ids) {
const int64_t num_reqs = query_end_loc.size(0);
const int64_t* target_ids_ptr = target_token_ids.data_ptr<int64_t>();
const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
const int32_t* query_end_ptr = query_end_loc.data_ptr<int32_t>();
int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
int64_t* out_pos_ptr = out_positions.data_ptr<int64_t>();
bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr<bool>();
bool* out_mask_ptr = out_is_masked_token_mask.data_ptr<bool>();
int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr<int32_t>();
int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr<int32_t>();
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
int32_t q_start = query_start_ptr[req_idx];
int32_t next_q_start = query_start_ptr[req_idx + 1];
int32_t q_end = query_end_ptr[req_idx];
int64_t num_valid_tokens =
shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1);
int64_t input_offset = shift_input_ids ? 1 : 0;
int64_t out_start = q_start + req_idx * (num_padding_slots_per_request -
(shift_input_ids ? 1 : 0));
int64_t num_rejected = next_q_start - q_end - 1;
int64_t total_output_tokens =
num_valid_tokens + num_padding_slots_per_request + num_rejected;
int64_t start_pos = target_pos_ptr[q_start];
int64_t bonus_token = next_ids_ptr[req_idx];
for (int64_t j = 0; j < total_output_tokens; ++j) {
int64_t out_idx = out_start + j;
bool is_valid = j < num_valid_tokens;
bool is_bonus = j == num_valid_tokens;
bool is_parallel = (j > num_valid_tokens) &&
(j < num_valid_tokens + num_padding_slots_per_request);
bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request;
int64_t in_idx =
std::min(static_cast<int64_t>(q_start + input_offset + j),
total_input_tokens - 1);
int64_t token_id = padding_token_id;
if (is_valid)
token_id = target_ids_ptr[in_idx];
else if (is_bonus)
token_id = bonus_token;
else if (is_parallel)
token_id = parallel_drafting_token_id;
out_ids_ptr[out_idx] = token_id;
out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j);
out_rej_mask_ptr[out_idx] = is_rejected;
out_mask_ptr[out_idx] = is_parallel;
if (is_bonus || is_parallel) {
int64_t new_token_local_idx = j - num_valid_tokens;
int64_t new_token_out_idx =
req_idx * num_padding_slots_per_request + new_token_local_idx;
out_new_idx_ptr[new_token_out_idx] = out_idx;
}
}
if (shift_input_ids) {
int64_t n_input = next_q_start - q_start;
for (int64_t j = 0; j < n_input; ++j) {
out_hidden_map_ptr[q_start + j] = out_start + j;
}
}
}
}
void rejection_greedy_sample_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
const torch::Tensor& bonus_token_ids,
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len) {
const int64_t batch_size = cu_num_draft_tokens.size(0);
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
const int64_t* target_argmax_ptr = target_argmax.data_ptr<int64_t>();
const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
const bool* greedy_ptr =
is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
const int64_t out_stride = output_token_ids.stride(0);
const int64_t bonus_stride = bonus_token_ids.stride(0);
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
if (greedy_ptr && !greedy_ptr[req_idx]) continue;
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
int64_t end_idx = cu_draft_ptr[req_idx];
int64_t num_draft_tokens = end_idx - start_idx;
bool rejected = false;
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
int64_t target_id = target_argmax_ptr[start_idx + pos];
out_ptr[req_idx * out_stride + pos] = target_id;
if (draft_ids_ptr[start_idx + pos] != target_id) {
rejected = true;
break;
}
}
if (!rejected) {
out_ptr[req_idx * out_stride + num_draft_tokens] =
bonus_ids_ptr[req_idx * bonus_stride];
}
}
}
void rejection_random_sample_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids,
const std::optional<torch::Tensor>& draft_probs,
const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
const torch::Tensor& recovered_token_ids,
const torch::Tensor& uniform_probs,
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
const int64_t vocab_size, const bool no_draft_probs) {
const int64_t batch_size = cu_num_draft_tokens.size(0);
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
const float* draft_probs_ptr =
no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
const float* target_probs_ptr = target_probs.data_ptr<float>();
const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr<int64_t>();
const float* uniform_probs_ptr = uniform_probs.data_ptr<float>();
const bool* greedy_ptr =
is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
const int64_t out_stride = output_token_ids.stride(0);
const int64_t bonus_stride = bonus_token_ids.stride(0);
const int64_t target_stride = target_probs.stride(0);
const int64_t draft_probs_stride =
no_draft_probs ? 0 : draft_probs.value().stride(0);
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
if (greedy_ptr && greedy_ptr[req_idx]) continue;
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
int64_t end_idx = cu_draft_ptr[req_idx];
int64_t num_draft_tokens = end_idx - start_idx;
bool rejected = false;
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
int64_t token_idx = start_idx + pos;
int64_t draft_id = draft_ids_ptr[token_idx];
float p = target_probs_ptr[token_idx * target_stride + draft_id];
float q =
no_draft_probs
? 1.0f
: draft_probs_ptr[token_idx * draft_probs_stride + draft_id];
float uniform_p = uniform_probs_ptr[token_idx];
float ratio = (q > 0.0f) ? (p / q) : 0.0f;
if (ratio >= uniform_p) {
out_ptr[req_idx * out_stride + pos] = draft_id;
} else {
out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx];
rejected = true;
break;
}
}
if (!rejected) {
out_ptr[req_idx * out_stride + num_draft_tokens] =
bonus_ids_ptr[req_idx * bonus_stride];
}
}
}
void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
const torch::Tensor& cu_num_tokens,
const int64_t replace_from, const int64_t replace_to) {
const int64_t batch_size = cu_num_tokens.size(0);
const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr<int64_t>();
int64_t* out_ptr = output.data_ptr<int64_t>();
const int64_t* in_ptr = input.data_ptr<int64_t>();
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1];
int64_t end_idx = cu_tokens_ptr[req_idx];
int64_t val = in_ptr[req_idx];
if (val == replace_from) {
val = replace_to;
}
for (int64_t i = start_idx; i < end_idx; ++i) {
out_ptr[i] = val;
}
}
}
void sample_recovered_tokens_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids,
const std::optional<torch::Tensor>& draft_probs,
const torch::Tensor& target_probs, const torch::Tensor& inv_q,
const int64_t vocab_size, const bool no_draft_probs) {
const int64_t batch_size = cu_num_draft_tokens.size(0);
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
const float* draft_probs_ptr =
no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
const float* target_probs_ptr = target_probs.data_ptr<float>();
const float* inv_q_ptr = inv_q.data_ptr<float>();
const int64_t target_stride = target_probs.stride(0);
const int64_t draft_probs_stride =
no_draft_probs ? 0 : draft_probs.value().stride(0);
const int64_t inv_q_stride = inv_q.stride(0);
#pragma omp parallel for
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
int64_t end_idx = cu_draft_ptr[req_idx];
int64_t num_draft_tokens = end_idx - start_idx;
const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride;
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
int64_t token_idx = start_idx + pos;
int64_t draft_id = draft_ids_ptr[token_idx];
const float* token_target_probs =
target_probs_ptr + token_idx * target_stride;
const float* token_draft_probs =
no_draft_probs ? nullptr
: (draft_probs_ptr + token_idx * draft_probs_stride);
int64_t best_id = 0;
float best_val = -1.0f;
for (int64_t v = 0; v < vocab_size; ++v) {
float prob = token_target_probs[v];
if (no_draft_probs) {
if (v == draft_id) prob = 0.0f;
} else {
float diff = prob - token_draft_probs[v];
prob = diff > 0.0f ? diff : 0.0f;
}
float val = prob * req_inv_q[v];
if (val > best_val) {
best_val = val;
best_id = v;
}
}
out_ptr[token_idx] = best_id;
}
}
}
} // namespace cpu_utils
-135
View File
@@ -85,9 +85,6 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
at::Tensor& w_scales,
std::optional<at::Tensor> bias);
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const std::string& activation);
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
@@ -141,63 +138,6 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& valid_sampled_tokens_count,
const torch::Tensor& query_start_loc_gpu,
torch::Tensor& token_indices_to_sample,
torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs);
void eagle_prepare_next_token_padded_kernel_impl(
const torch::Tensor& sampled_token_ids,
const torch::Tensor& discard_request_mask,
const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
const int64_t num_sampled_tokens_per_req, const int64_t num_reqs);
void eagle_step_slot_mapping_metadata_kernel_impl(
const torch::Tensor& positions, const torch::Tensor& block_table,
torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
torch::Tensor& out_slot_mapping, const int64_t block_size,
const int64_t max_model_len, const int64_t PAD_ID);
void copy_and_expand_eagle_inputs_kernel_impl(
const torch::Tensor& target_token_ids,
const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
torch::Tensor& out_input_ids, torch::Tensor& out_positions,
torch::Tensor& out_is_rejected_token_mask,
torch::Tensor& out_is_masked_token_mask,
torch::Tensor& out_new_token_indices,
torch::Tensor& out_hidden_state_mapping,
const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
const int64_t total_input_tokens,
const int64_t num_padding_slots_per_request, const bool shift_input_ids);
void rejection_greedy_sample_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
const torch::Tensor& bonus_token_ids,
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len);
void rejection_random_sample_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids,
const std::optional<torch::Tensor>& draft_probs,
const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
const torch::Tensor& recovered_token_ids,
const torch::Tensor& uniform_probs,
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
const int64_t vocab_size, const bool no_draft_probs);
void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
const torch::Tensor& cu_num_tokens,
const int64_t replace_from, const int64_t replace_to);
void sample_recovered_tokens_kernel_impl(
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
const torch::Tensor& draft_token_ids,
const std::optional<torch::Tensor>& draft_probs,
const torch::Tensor& target_probs, const torch::Tensor& inv_q,
const int64_t vocab_size, const bool no_draft_probs);
} // namespace cpu_utils
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
@@ -236,15 +176,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
#if (defined(__aarch64__) && !defined(__APPLE__))
ops.def(
"activation_lut_bf16(Tensor! out, Tensor input, str activation)"
" -> ()");
ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
#endif // (defined(__aarch64__) && !defined(__APPLE__))
// Layernorm
// Apply Root Mean Square (RMS) Normalization to the input tensor.
ops.def(
@@ -432,72 +363,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt "
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
ops.def(
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
"Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, "
"Tensor(a3!) token_indices_to_sample, "
"Tensor(a4!) num_rejected_tokens_gpu, "
"SymInt num_reqs) -> ()",
&cpu_utils::eagle_prepare_inputs_padded_kernel_impl);
ops.def(
"eagle_prepare_next_token_padded_kernel_impl("
"Tensor sampled_token_ids, Tensor discard_request_mask, "
"Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, "
"Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, "
"SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()",
&cpu_utils::eagle_prepare_next_token_padded_kernel_impl);
ops.def(
"eagle_step_slot_mapping_metadata_kernel_impl("
"Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, "
"Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, "
"SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()",
&cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl);
ops.def(
"copy_and_expand_eagle_inputs_kernel_impl("
"Tensor target_token_ids, Tensor target_positions, "
"Tensor next_token_ids, Tensor(a3!) out_input_ids, "
"Tensor(a4!) out_positions, "
"Tensor(a5!) out_is_rejected_token_mask, "
"Tensor(a6!) out_is_masked_token_mask, "
"Tensor(a7!) out_new_token_indices, "
"Tensor(a8!) out_hidden_state_mapping, "
"Tensor query_start_loc, Tensor query_end_loc, "
"SymInt padding_token_id, SymInt parallel_drafting_token_id, "
"SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
"bool shift_input_ids) -> ()",
&cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
ops.def(
"rejection_greedy_sample_kernel_impl("
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
"Tensor draft_token_ids, Tensor target_argmax, "
"Tensor bonus_token_ids, Tensor? is_greedy, "
"SymInt max_spec_len) -> ()",
&cpu_utils::rejection_greedy_sample_kernel_impl);
ops.def(
"rejection_random_sample_kernel_impl("
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
"Tensor draft_token_ids, Tensor? draft_probs, "
"Tensor target_probs, Tensor bonus_token_ids, "
"Tensor recovered_token_ids, Tensor uniform_probs, "
"Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, "
"bool no_draft_probs) -> ()",
&cpu_utils::rejection_random_sample_kernel_impl);
ops.def(
"expand_kernel_impl(Tensor(a0!) output, Tensor input, "
"Tensor cu_num_tokens, SymInt replace_from, "
"SymInt replace_to) -> ()",
&cpu_utils::expand_kernel_impl);
ops.def(
"sample_recovered_tokens_kernel_impl("
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
"Tensor draft_token_ids, Tensor? draft_probs, "
"Tensor target_probs, Tensor inv_q, SymInt vocab_size, "
"bool no_draft_probs) -> ()",
&cpu_utils::sample_recovered_tokens_kernel_impl);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
+6 -73
View File
@@ -13,80 +13,13 @@
#include "cpu/utils.hpp"
#ifdef VLLM_NUMA_DISABLED
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
#else
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
// Memory node binding
if (numa_available() != -1) {
// Concatenate all node_ids into a single comma-separated string
if (!node_ids.empty()) {
std::string node_ids_str;
for (const int node_id : node_ids) {
if (!node_ids_str.empty()) {
node_ids_str += ",";
}
node_ids_str += std::to_string(node_id);
}
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
bitmask* src_mask = numa_get_mems_allowed();
int pid = getpid();
if (mask && src_mask) {
// move all existing pages to the specified numa node.
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
int page_num = numa_migrate_pages(pid, src_mask, mask);
if (page_num == -1) {
TORCH_WARN("numa_migrate_pages failed. errno: " +
std::to_string(errno));
}
// Restrict memory allocation to the selected NUMA node(s).
// Enhances memory locality for the threads bound to those NUMA CPUs.
if (node_ids.size() > 1) {
errno = 0;
numa_set_interleave_mask(mask);
if (errno != 0) {
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using INTERLEAVE policy for memory "
"allocation across multiple NUMA nodes (nodes: " +
node_ids_str +
"). Memory allocations will be "
"interleaved across the specified NUMA nodes.");
}
} else {
errno = 0;
numa_set_membind(mask);
if (errno != 0) {
TORCH_WARN("numa_set_membind failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using MEMBIND policy for memory "
"allocation on the NUMA nodes (" +
node_ids_str +
"). Memory allocations will be "
"strictly bound to these NUMA nodes.");
}
}
numa_set_strict(1);
numa_free_nodemask(mask);
numa_free_nodemask(src_mask);
} else {
TORCH_WARN(
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
std::to_string(errno));
}
}
}
std::string init_cpu_threads_env(const std::string& cpu_ids) {
return std::string(
"Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
"no effect to setup thread affinity.");
}
#endif // VLLM_NUMA_DISABLED
#endif
namespace cpu_utils {
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
+1 -24
View File
@@ -54,34 +54,11 @@ struct Counter {
};
inline int64_t get_available_l2_size() {
#if defined(__s390x__)
static int64_t size = []() {
uint32_t l2_cache_size = 0;
auto caps = at::cpu::get_cpu_capabilities();
auto it = caps.find("l2_cache_size");
if (it != caps.end()) {
l2_cache_size = static_cast<uint32_t>(it->second.toInt());
}
if (l2_cache_size == 0) {
long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
if (sys_l2 > 0) {
l2_cache_size = static_cast<uint32_t>(sys_l2);
}
}
if (l2_cache_size == 0) {
l2_cache_size = 256 * 1024;
}
return static_cast<int64_t>(l2_cache_size) >> 1; // use 50% of L2 cache
}();
return size;
#else
static int64_t size = []() {
auto caps = at::cpu::get_cpu_capabilities();
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
const uint32_t l2_cache_size = at::cpu::L2_cache_size();
return l2_cache_size >> 1; // use 50% of L2 cache
}();
return size;
#endif
}
template <int32_t alignment_v, typename T>
@@ -389,28 +389,20 @@ struct Sm90ColOrScalarBroadcastArray {
CUTLASS_DEVICE void
begin() {
cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol(i)) < m;
}
if (!params.col_broadcast) {
fill(tCrCol, *(params.ptr_col_array[group]));
return;
}
// tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and
// EPI_N are stride-0 for the column broadcast. Slice those modes at
// index 0 to avoid redundant copies AND ensure pred/data consistency
static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast");
static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast");
auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
cute::Tensor pred = make_tensor<bool>(shape(tCgCol_s));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol_s(i)) < m;
}
copy_if(pred, tCgCol_s, tCrCol_s);
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_if(pred, filter(tCgCol), filter(tCrCol));
}
template <typename ElementAccumulator, int FragmentSize>
@@ -382,28 +382,20 @@ struct Sm90ColOrScalarBroadcast {
CUTLASS_DEVICE void
begin() {
cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol(i)) < m;
}
if (!params.col_broadcast) {
fill(tCrCol, *(params.ptr_col));
return;
}
// tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and
// EPI_N are stride-0 for the column broadcast. Slice those modes at
// index 0 to avoid redundant copies AND ensure pred/data consistency
static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast");
static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast");
auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M)
cute::Tensor pred = make_tensor<bool>(shape(tCgCol_s));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol_s(i)) < m;
}
copy_if(pred, tCgCol_s, tCrCol_s);
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_if(pred, filter(tCgCol), filter(tCrCol));
}
template <typename ElementAccumulator, int FragmentSize>
+9 -388
View File
@@ -19,10 +19,8 @@
#include <type_traits>
#include <torch/cuda.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "async_util.cuh"
#include "cuda_compat.h"
#include "dispatch_utils.h"
#include "type_convert.cuh"
@@ -88,9 +86,6 @@ inline __device__ __host__ T divUp(T m, T n) {
} // namespace tensorrt_llm::common
namespace tensorrt_llm::kernels {
using namespace vllm::cuda_async;
// NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation,
// with added support for passing the cos_sin_cache as an input.
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
@@ -306,237 +301,6 @@ __global__ void fusedQKNormRopeKernel(
#endif
}
// Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for
// the same token, sharing cos/sin from shared memory via cp.async.
// When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads,
// hiding global-memory latency and improving occupancy for large batches.
template <typename scalar_t_in, typename scalar_t_cache, int head_dim,
bool interleave, int HEADS_PER_WARP>
__global__ void fusedQKNormRopeKernelNTokenHeads(
void* qkv_void, int const num_heads_q, int const num_heads_k,
int const num_heads_v, float const eps, void const* q_weight_void,
void const* k_weight_void, void const* cos_sin_cache_void,
int64_t const* position_ids, int const num_tokens, int const rotary_dim) {
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
std::is_same_v<scalar_t_cache, c10::BFloat16>) {
return;
} else {
#endif
using Converter = vllm::_typeConvert<scalar_t_in>;
static_assert(Converter::exists,
"Input QKV data type is not supported for this CUDA "
"architecture or toolkit version.");
using T_in = typename Converter::hip_type;
using T2_in = typename Converter::packed_hip_type;
using CacheConverter = vllm::_typeConvert<scalar_t_cache>;
static_assert(CacheConverter::exists,
"Cache data type is not supported for this CUDA architecture "
"or toolkit version.");
using T_cache = typename CacheConverter::hip_type;
extern __shared__ char smem_storage[];
// Shared memory layout:
// [0, cos_sin_bytes) : cos/sin for each warp (warpsPerBlock *
// rotary_dim * sizeof(T_cache))
// [cos_sin_bytes, ...) : QKV tiles
// per warp (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes)
T_cache* const smem = reinterpret_cast<T_cache*>(smem_storage);
T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
T_cache const* cos_sin_cache =
reinterpret_cast<T_cache const*>(cos_sin_cache_void);
int const warpsPerBlock = blockDim.x / 32;
int const warpId = threadIdx.x / 32;
int const laneId = threadIdx.x % 32;
int const total_qk_heads = num_heads_q + num_heads_k;
int const num_heads = num_heads_q + num_heads_k + num_heads_v;
int const head_chunks_per_token =
(total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP;
int const warp_global = blockIdx.x * warpsPerBlock + warpId;
int const tokenIdx = warp_global / head_chunks_per_token;
int const headChunk = warp_global % head_chunks_per_token;
int const first_head = headChunk * HEADS_PER_WARP;
int const num_heads_this_warp =
(first_head + HEADS_PER_WARP <= total_qk_heads)
? HEADS_PER_WARP
: (total_qk_heads - first_head);
if (tokenIdx >= num_tokens) return;
static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64");
constexpr int numElemsPerThread = head_dim / 32;
constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
static_assert(elemSizeBytes % 4 == 0,
"elemSizeBytes must be a multiple of 4");
constexpr int vecSize = elemSizeBytes / 4;
using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type;
int const cos_sin_bytes =
warpsPerBlock * rotary_dim * static_cast<int>(sizeof(T_cache));
int const qkv_tile_bytes = 32 * elemSizeBytes;
char* const this_warp_head_smem =
smem_storage + cos_sin_bytes +
warpId * (HEADS_PER_WARP * qkv_tile_bytes);
// === Group 0: async load all heads' QKV into smem (issued first). ===
for (int k = 0; k < num_heads_this_warp; ++k) {
int const localHeadIdx = first_head + k;
bool const isQ = localHeadIdx < num_heads_q;
int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
int offWarp;
if (isQ) {
offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
} else {
offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
headIdx * head_dim;
}
int const offThread = offWarp + laneId * numElemsPerThread;
char* smem_dst =
this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
cp_async_shared_global_ca(smem_dst,
reinterpret_cast<const char*>(&qkv[offThread]),
elemSizeBytes);
}
cp_async_commit_group(); // commit group 0 (QKV)
// === Group 1: async load cos/sin into smem (issued second). ===
int64_t const pos_id = position_ids[tokenIdx];
T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim;
int const copy_bytes = rotary_dim * static_cast<int>(sizeof(T_cache));
int const num_copies = (copy_bytes + 15) / 16;
for (int copyId = laneId; copyId < num_copies; copyId += 32) {
char* smem_ptr =
reinterpret_cast<char*>(&smem[warpId * rotary_dim]) + copyId * 16;
const char* glob_ptr =
reinterpret_cast<const char*>(cache_ptr) + copyId * 16;
cp_async_shared_global_16_cg(smem_ptr, glob_ptr);
}
cp_async_commit_group(); // commit group 1 (cos/sin)
// wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done.
cp_async_wait_group<1>();
float elements[numElemsPerThread];
float elements2[numElemsPerThread];
int const rotary_lanes = rotary_dim / numElemsPerThread;
int const embed_dim = rotary_dim / 2;
T_cache const* const cos_smem = &smem[warpId * rotary_dim];
T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim];
// Preload weights into registers once, reused across all heads.
float q_w[numElemsPerThread];
float k_w[numElemsPerThread];
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
int const dim = laneId * numElemsPerThread + i;
q_w[i] = Converter::convert(q_weight[dim]);
k_w[i] = Converter::convert(k_weight[dim]);
}
for (int k = 0; k < num_heads_this_warp; ++k) {
int const localHeadIdx = first_head + k;
bool const isQ = localHeadIdx < num_heads_q;
int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
int offsetWarp;
if (isQ) {
offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
} else {
offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
headIdx * head_dim;
}
int const offsetThread = offsetWarp + laneId * numElemsPerThread;
// === Part 1: QK Norm (read from smem; group 0 already done). ===
float sumOfSquares = 0.0f;
{
char const* smem_src =
this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
vec_T vec = *reinterpret_cast<vec_T const*>(smem_src);
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
#pragma unroll
for (int i = 0; i < num_packed_elems; i++) {
T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
float2 vals = Converter::convert(packed_val);
sumOfSquares += vals.x * vals.x;
sumOfSquares += vals.y * vals.y;
elements[2 * i] = vals.x;
elements[2 * i + 1] = vals.y;
}
}
sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares);
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]);
}
// On first head: wait for group 1 (cos/sin) before RoPE.
if (k == 0) cp_async_wait_group<0>();
// === Part 2: RoPE using cos/sin from shared memory. ===
if (laneId < rotary_lanes) {
if constexpr (interleave) {
#pragma unroll
for (int i = 0; i < numElemsPerThread / 2; ++i) {
int const idx0 = 2 * i;
int const idx1 = 2 * i + 1;
int const dim_idx = laneId * numElemsPerThread + idx0;
float const val0 = elements[idx0];
float const val1 = elements[idx1];
int const half_dim = dim_idx / 2;
float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
elements[idx0] = val0 * cos_val - val1 * sin_val;
elements[idx1] = val0 * sin_val + val1 * cos_val;
}
} else {
__syncwarp();
int const pairOffset = (rotary_dim / 2) / numElemsPerThread;
#pragma unroll
for (int i = 0; i < numElemsPerThread; i++) {
elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset);
if (laneId < pairOffset) elements2[i] = -elements2[i];
int dim_idx = laneId * numElemsPerThread + i;
dim_idx = (dim_idx * 2) % rotary_dim;
int const half_dim = dim_idx / 2;
float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
}
__syncwarp();
}
}
// Store.
{
vec_T vec;
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
#pragma unroll
for (int i = 0; i < num_packed_elems; i++) {
T2_in packed_val = Converter::convert(
make_float2(elements[2 * i], elements[2 * i + 1]));
*(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
}
*reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
}
}
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
}
#endif
}
// Borrowed from
// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
@@ -557,12 +321,15 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
void const* cos_sin_cache, bool const interleave,
int64_t const* position_ids, cudaStream_t stream) {
constexpr int blockSize = 256;
int const warpsPerBlock = blockSize / 32;
int const totalQKHeads = num_heads_q + num_heads_k;
int const totalWarps = num_tokens * totalQKHeads;
int const gridSize = common::divUp(totalWarps, warpsPerBlock);
dim3 gridDim(gridSize);
dim3 blockDim(blockSize);
switch (head_dim) {
case 64:
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
@@ -593,118 +360,6 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
"Unsupported head dimension for fusedQKNormRope: ", head_dim);
}
}
// Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8).
// When token_heads_per_warp == 1, delegates to the 1-head baseline above.
template <typename scalar_t_in, typename scalar_t_cache>
void launchFusedQKNormRopeNTokenHeads(
void* qkv, int const num_tokens, int const num_heads_q,
int const num_heads_k, int const num_heads_v, int const head_dim,
int const rotary_dim, float const eps, void const* q_weight,
void const* k_weight, void const* cos_sin_cache, bool const interleave,
int64_t const* position_ids, int const token_heads_per_warp,
cudaStream_t stream) {
TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
token_heads_per_warp == 4 || token_heads_per_warp == 8,
"token_heads_per_warp must be 1, 2, 4, or 8, got ",
token_heads_per_warp);
// token_heads_per_warp == 1: delegate to the 1-head baseline kernel.
if (token_heads_per_warp == 1) {
launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
position_ids, stream);
return;
}
// NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks.
// If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last
// cp.async would write past the shared memory allocation.
// Fall back to the base kernel instead of failing.
{
size_t const rotary_bytes =
static_cast<size_t>(rotary_dim) *
(std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u);
if (rotary_bytes % 16 != 0) {
launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
position_ids, stream);
return;
}
}
constexpr int blockSize = 256;
int const warpsPerBlock = blockSize / 32;
int const totalQKHeads = num_heads_q + num_heads_k;
// Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem.
int const head_chunks_per_token =
(totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp;
int const total_warps = num_tokens * head_chunks_per_token;
int const gridSize = common::divUp(total_warps, warpsPerBlock);
dim3 gridDim(gridSize);
dim3 blockDim(blockSize);
// Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same
// layout).
size_t const cache_elem_size =
std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u;
// QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2)
// = 2*head_dim bytes.
size_t const qkv_smem_per_warp = static_cast<size_t>(token_heads_per_warp) *
2u * static_cast<size_t>(head_dim);
size_t const smem_bytes =
warpsPerBlock * static_cast<size_t>(rotary_dim) * cache_elem_size +
warpsPerBlock * qkv_smem_per_warp;
#define LAUNCH_N_TOKEN_HEADS(N) \
do { \
switch (head_dim) { \
case 64: \
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 64, \
INTERLEAVE, (N)> \
<<<gridDim, blockDim, smem_bytes, stream>>>( \
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
k_weight, cos_sin_cache, position_ids, num_tokens, \
rotary_dim); \
}); \
break; \
case 128: \
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 128, \
INTERLEAVE, (N)> \
<<<gridDim, blockDim, smem_bytes, stream>>>( \
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
k_weight, cos_sin_cache, position_ids, num_tokens, \
rotary_dim); \
}); \
break; \
case 256: \
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 256, \
INTERLEAVE, (N)> \
<<<gridDim, blockDim, smem_bytes, stream>>>( \
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
k_weight, cos_sin_cache, position_ids, num_tokens, \
rotary_dim); \
}); \
break; \
default: \
TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
} \
} while (0)
if (token_heads_per_warp == 2) {
LAUNCH_N_TOKEN_HEADS(2);
} else if (token_heads_per_warp == 4) {
LAUNCH_N_TOKEN_HEADS(4);
} else if (token_heads_per_warp == 8) {
LAUNCH_N_TOKEN_HEADS(8);
}
#undef LAUNCH_N_TOKEN_HEADS
}
} // namespace tensorrt_llm::kernels
void fused_qk_norm_rope(
@@ -719,8 +374,7 @@ void fused_qk_norm_rope(
torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim]
bool is_neox, // Whether RoPE is applied in Neox style
torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens]
int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value
torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
) {
// Input validation
CHECK_INPUT(qkv);
@@ -760,48 +414,15 @@ void fused_qk_norm_rope(
qkv.size(1) == total_heads * head_dim,
"QKV tensor size must match total number of heads and head dimension");
auto device_id = qkv.get_device();
auto stream = at::cuda::getCurrentCUDAStream(device_id);
// Select token_heads_per_warp: forced value if >0, else auto-select.
// Auto thresholds are calibrated on SM 9.0 (H100). On other architectures,
// fall back to token_heads_per_warp=1 (base kernel) until profiled.
int token_heads_per_warp;
if (forced_token_heads_per_warp > 0) { // only support SM80+
token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp);
} else {
token_heads_per_warp = 1;
auto* dev_prop = at::cuda::getDeviceProperties(device_id);
int sm_version = dev_prop->major * 10 + dev_prop->minor;
int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k);
if (sm_version == 90) {
if (head_dim >= 256) {
if (total_qk_units < 4096LL) {
token_heads_per_warp = 1;
} else if (total_qk_units < 8192LL) {
token_heads_per_warp = 2;
} else {
token_heads_per_warp = 4;
}
} else {
if (total_qk_units < 10240LL) {
token_heads_per_warp = 1;
} else if (total_qk_units < 40960LL) {
token_heads_per_warp = 4;
} else {
token_heads_per_warp = 8;
}
}
}
}
auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
using qkv_scalar_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(
cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
using cache_scalar_t = scalar_t;
tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
qkv_scalar_t, cache_scalar_t>(
tensorrt_llm::kernels::launchFusedQKNormRope<qkv_scalar_t,
cache_scalar_t>(
qkv.data_ptr(), static_cast<int>(num_tokens),
static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
static_cast<int>(num_heads_v), static_cast<int>(head_dim),
@@ -809,7 +430,7 @@ void fused_qk_norm_rope(
q_weight.data_ptr(), k_weight.data_ptr(),
cos_sin_cache.data_ptr(), !is_neox,
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
token_heads_per_warp, stream);
stream);
});
});
}
}
-9
View File
@@ -134,13 +134,4 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& input_global_scale);
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets);
#endif

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