forked from Karylab-cklius/vllm
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b547993038 | ||
|
|
9026d5ec05 | ||
|
|
ee778e2cd8 |
@@ -46,7 +46,7 @@ echo "Image not found, proceeding with build..."
|
||||
|
||||
# --- CUDA 13.0 for nightly builds ---
|
||||
# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
|
||||
NIGHTLY_CUDA_VERSION="13.0.2"
|
||||
NIGHTLY_CUDA_VERSION="13.0.0"
|
||||
NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
|
||||
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
'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_functions.py &&
|
||||
pytest -v -s lora/test_lora_huggingface.py &&
|
||||
pytest -v -s lora/test_lora_manager.py &&
|
||||
pytest -v -s lora/test_lora_utils.py &&
|
||||
@@ -125,6 +125,6 @@ steps:
|
||||
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_qwen3_unembed.py &&
|
||||
pytest -v -s lora/test_qwenvl.py &&
|
||||
pytest -v -s lora/test_whisper.py'
|
||||
|
||||
@@ -1,13 +1,3 @@
|
||||
# CUDA architecture lists — following PyTorch RELEASE.md
|
||||
# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
|
||||
# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
|
||||
env:
|
||||
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
# aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
|
||||
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
|
||||
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
|
||||
|
||||
steps:
|
||||
- input: "Provide Release version here"
|
||||
id: input-release-version
|
||||
@@ -24,10 +14,12 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -37,7 +29,9 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -63,7 +57,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
@@ -76,7 +70,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -114,95 +108,96 @@ steps:
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-12-9
|
||||
id: build-release-image-x86-cuda-13-0
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-12-9
|
||||
id: build-release-image-arm64-cuda-13-0
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
id: build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
id: build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
key: block-cpu-release-image-build
|
||||
@@ -243,7 +238,7 @@ steps:
|
||||
- group: "Publish release images"
|
||||
key: "publish-release-images"
|
||||
steps:
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
depends_on:
|
||||
- build-release-image-x86
|
||||
- build-release-image-arm64
|
||||
@@ -255,7 +250,7 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Annotate release workflow - CUDA 13.0"
|
||||
- label: "Annotate release workflow - CUDA 12.9"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
@@ -264,19 +259,19 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-12-9
|
||||
- build-release-image-arm64-cuda-12-9
|
||||
id: create-multi-arch-manifest-cuda-12-9
|
||||
- build-release-image-x86-cuda-13-0
|
||||
- build-release-image-arm64-cuda-13-0
|
||||
id: create-multi-arch-manifest-cuda-13-0
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-ubuntu2404
|
||||
- build-release-image-arm64-ubuntu2404
|
||||
@@ -288,17 +283,17 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
- build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
|
||||
- build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
- build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
@@ -318,16 +313,16 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest-cuda-12-9
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu129"
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
|
||||
@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
|
||||
(Optional) For CUDA 12.9:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
(Optional) For CUDA 13.0:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
@@ -33,8 +33,8 @@ To download and upload the image:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
|
||||
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
|
||||
## ROCm
|
||||
|
||||
@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129
|
||||
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker manifest push vllm/vllm-openai:latest-cu129
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
|
||||
docker manifest rm vllm/vllm-openai:latest-cu130
|
||||
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker manifest push vllm/vllm-openai:latest-cu130
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
|
||||
|
||||
docker manifest rm vllm/vllm-openai-cpu:latest || true
|
||||
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
|
||||
|
||||
@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
|
||||
TORCH_INDEX_URL=""
|
||||
fi
|
||||
else
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
|
||||
fi
|
||||
echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ set -ex
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
INDICES_OUTPUT_DIR="indices"
|
||||
DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
@@ -25,100 +25,22 @@ export PYTHONPATH=".."
|
||||
###############################################################################
|
||||
|
||||
cleanup_docker() {
|
||||
# Share the same lock with image pull to avoid cleanup/pull races on one node.
|
||||
local docker_lock="/tmp/docker-pull.lock"
|
||||
exec 9>"$docker_lock"
|
||||
flock 9
|
||||
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
if [ -z "$docker_root" ]; then
|
||||
echo "Failed to determine Docker root directory." >&2
|
||||
flock -u 9
|
||||
return 1
|
||||
exit 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
docker image prune -f
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
else
|
||||
echo "Disk usage is below $threshold%. Checking old CI images anyway."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
|
||||
fi
|
||||
echo "Old CI image cleanup completed."
|
||||
|
||||
flock -u 9
|
||||
}
|
||||
|
||||
cleanup_old_ci_images() {
|
||||
local repo_prefix="$1"
|
||||
local current_image_ref="$2"
|
||||
local ttl_hours="$3"
|
||||
local aggressive_cleanup="$4"
|
||||
|
||||
if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
|
||||
echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
|
||||
echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
|
||||
ttl_hours=72
|
||||
fi
|
||||
|
||||
local now_epoch cutoff_epoch
|
||||
now_epoch=$(date +%s)
|
||||
cutoff_epoch=$((now_epoch - ttl_hours * 3600))
|
||||
|
||||
local -a used_image_ids
|
||||
mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
|
||||
|
||||
local removed_count=0
|
||||
local examined_count=0
|
||||
declare -A seen_ids=()
|
||||
|
||||
while read -r image_ref image_id; do
|
||||
[[ -z "$image_ref" || -z "$image_id" ]] && continue
|
||||
((examined_count++))
|
||||
|
||||
# Keep the image this job is going to use.
|
||||
if [[ "$image_ref" == "$current_image_ref" ]]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Avoid duplicate deletes when multiple tags point to same image id.
|
||||
if [[ -n "${seen_ids[$image_id]:-}" ]]; then
|
||||
continue
|
||||
fi
|
||||
seen_ids[$image_id]=1
|
||||
|
||||
# Never delete images that are used by any container on this node.
|
||||
if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
local created created_epoch
|
||||
created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
|
||||
[[ -z "$created" ]] && continue
|
||||
created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
|
||||
[[ -z "$created_epoch" ]] && continue
|
||||
|
||||
if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
if docker image rm -f "$image_id" >/dev/null 2>&1; then
|
||||
((removed_count++))
|
||||
fi
|
||||
fi
|
||||
done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
|
||||
|
||||
# Also trim old dangling layers; this is safe and does not remove referenced images.
|
||||
docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
|
||||
|
||||
if [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
|
||||
else
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
|
||||
echo "Disk usage is below $threshold%. No cleanup needed."
|
||||
fi
|
||||
}
|
||||
|
||||
@@ -343,6 +265,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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+14
-14
@@ -388,10 +388,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1647,10 +1647,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1951,8 +1951,8 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
@@ -2930,10 +2930,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
|
||||
@@ -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
|
||||
@@ -226,6 +226,91 @@ steps:
|
||||
commands:
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -95,13 +95,11 @@ steps:
|
||||
- tests/kernels/moe/test_deepgemm.py
|
||||
- tests/kernels/moe/test_batched_deepgemm.py
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
- pytest -v -s kernels/attention/test_deepgemm_attention.py
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -113,10 +113,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# for multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -44,10 +44,10 @@ steps:
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -69,9 +69,9 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- 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/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
- 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/"
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- 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/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- 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/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
|
||||
@@ -28,7 +28,6 @@ steps:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
|
||||
+6
-2
@@ -262,7 +262,7 @@ pull_request_rules:
|
||||
- files~=^docker/Dockerfile.xpu
|
||||
- files~=^\\.buildkite/intel_jobs/
|
||||
- files=\.buildkite/ci_config_intel.yaml
|
||||
- files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
|
||||
- files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
|
||||
- files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
|
||||
@@ -389,7 +389,11 @@ pull_request_rules:
|
||||
- files~=^tests/entrypoints/anthropic/.*tool.*
|
||||
- files~=^vllm/tool_parsers/
|
||||
- files=docs/features/tool_calling.md
|
||||
- files~=^examples/tool_calling/
|
||||
- files~=^examples/tool_chat_*
|
||||
- files=examples/offline_inference/chat_with_tools.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
|
||||
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
|
||||
@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
|
||||
# Limit the number of parallel jobs to avoid OOM
|
||||
export MAX_JOBS=1
|
||||
# Make sure release wheels are built for the following architectures
|
||||
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
|
||||
|
||||
bash tools/check_repo.sh
|
||||
|
||||
|
||||
+7
-13
@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
|
||||
# Supported python versions. These versions will be searched in order, the
|
||||
# first match will be selected. These should be kept in sync with setup.py.
|
||||
#
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
|
||||
|
||||
# Supported AMD GPU architectures.
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
|
||||
# ROCm installation prefix. Default to /opt/rocm but allow override via
|
||||
# -DROCM_PATH=/your/rocm/path when invoking cmake.
|
||||
@@ -94,15 +94,12 @@ find_package(Torch REQUIRED)
|
||||
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
|
||||
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
|
||||
# to support the whole generation without specifying all sub-architectures
|
||||
# see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
|
||||
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
|
||||
else()
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
endif()
|
||||
|
||||
#
|
||||
@@ -310,9 +307,7 @@ set(VLLM_EXT_SRC
|
||||
"csrc/torch_bindings.cpp")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/minimax_reduce_rms_kernel.cu"
|
||||
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
|
||||
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
@@ -1053,8 +1048,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_MOE_EXT_SRC
|
||||
"csrc/moe/moe_wna16.cu"
|
||||
"csrc/moe/grouped_topk_kernels.cu"
|
||||
"csrc/moe/router_gemm.cu"
|
||||
"csrc/moe/topk_softplus_sqrt_kernels.cu")
|
||||
"csrc/moe/router_gemm.cu")
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
|
||||
| <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
|
||||
</p>
|
||||
|
||||
🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
|
||||
|
||||
---
|
||||
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
|
||||
- Efficient multi-LoRA support for dense and MoE layers
|
||||
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||
|
||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||
vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
|
||||
|
||||
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
|
||||
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
|
||||
|
||||
@@ -404,7 +404,6 @@ def _build_attention_metadata(
|
||||
query_start_loc=q_start_gpu,
|
||||
query_start_loc_cpu=q_start_cpu,
|
||||
seq_lens=seq_lens_gpu,
|
||||
seq_lens_cpu_upper_bound=seq_lens_cpu,
|
||||
_seq_lens_cpu=seq_lens_cpu,
|
||||
_num_computed_tokens_cpu=num_computed_tokens_cpu,
|
||||
slot_mapping=slot_mapping,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -1,324 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
|
||||
# python3 benchmark_vit_fp8_attn.py
|
||||
#
|
||||
# Profile mode (PyTorch profiler, saves TensorBoard traces):
|
||||
# python3 benchmark_vit_fp8_attn.py --profile
|
||||
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
|
||||
#
|
||||
# Custom seq_lens:
|
||||
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
|
||||
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.profiler import ProfilerActivity, profile, record_function
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Qwen3-VL defaults
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
|
||||
|
||||
|
||||
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
|
||||
"""Create FP8 and BF16 attention modules + workspace."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
backend_patch = patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
)
|
||||
|
||||
# FP8 attention
|
||||
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config_fp8 = VllmConfig()
|
||||
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
|
||||
with set_current_vllm_config(vllm_config_fp8), backend_patch:
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
# BF16 attention (no FP8)
|
||||
with set_current_vllm_config(VllmConfig()), backend_patch:
|
||||
attn_bf16 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
workspace = _get_flashinfer_workspace_buffer()
|
||||
return attn_fp8, attn_bf16, workspace
|
||||
|
||||
|
||||
def _build_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8: bool,
|
||||
):
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
cu_np = np.array([0, seq_len], dtype=np.int32)
|
||||
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
|
||||
|
||||
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
|
||||
)
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_np,
|
||||
num_heads * head_dim,
|
||||
1,
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded,
|
||||
)
|
||||
return cu_seqlens, max_seqlen, seq_lengths
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
seq_lens: list[int],
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
method: str,
|
||||
):
|
||||
"""Benchmark FP8 vs BF16 attention across seq_lens.
|
||||
|
||||
Uses FlashInfer GPU-level timing to measure pure kernel time,
|
||||
excluding CPU launch overhead.
|
||||
"""
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
|
||||
|
||||
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
|
||||
elif method == "cudagraph":
|
||||
from flashinfer.testing import (
|
||||
bench_gpu_time_with_cudagraph as bench_fn,
|
||||
)
|
||||
|
||||
bench_fn = partial(bench_fn, cold_l2_cache=False)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
print(f"Timing method: {method}")
|
||||
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
|
||||
print("-" * 46)
|
||||
|
||||
for seq_len in seq_lens:
|
||||
torch.manual_seed(42)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
|
||||
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
|
||||
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
# bench_fn returns List[float] of per-iteration times in ms
|
||||
bf16_times = bench_fn(bf16_fn)
|
||||
fp8_times = bench_fn(fp8_fn)
|
||||
|
||||
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
|
||||
fp8_us = np.median(fp8_times) * 1e3
|
||||
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
|
||||
|
||||
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
|
||||
|
||||
|
||||
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
|
||||
"""Create a trace handler that saves to TensorBoard and prints summary."""
|
||||
|
||||
def handler(prof):
|
||||
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
|
||||
print(f"\n{'=' * 80}")
|
||||
print(label)
|
||||
print(f"{'=' * 80}")
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def run_profile(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
warmup: int,
|
||||
output_dir: str,
|
||||
):
|
||||
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
|
||||
|
||||
# Profile BF16 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"bf16_h{head_dim}_s{seq_len}",
|
||||
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_bf16:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("bf16_attention"):
|
||||
attn_bf16._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_bf16.step()
|
||||
|
||||
# Profile FP8 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"fp8_h{head_dim}_s{seq_len}",
|
||||
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_fp8:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("fp8_attention"):
|
||||
attn_fp8._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_fp8.step()
|
||||
|
||||
print(f"\nTensorBoard traces saved to: {output_dir}")
|
||||
print(f"View with: tensorboard --logdir={output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=DEFAULT_SEQ_LENS,
|
||||
help="Sequence lengths to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
default=NUM_HEADS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head-dim",
|
||||
type=int,
|
||||
default=HEAD_DIM,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["cupti", "cudagraph"],
|
||||
default="cudagraph",
|
||||
help="GPU timing method: cupti (CUPTI kernel timing) or "
|
||||
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (profile mode only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help="Run PyTorch profiler instead of benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-seq-len",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Sequence length for profiling (default: 8192)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-output-dir",
|
||||
type=str,
|
||||
default="./profile_traces",
|
||||
help="Output directory for TensorBoard traces (default: ./profile_traces)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.profile:
|
||||
run_profile(
|
||||
args.profile_seq_len,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.warmup,
|
||||
args.profile_output_dir,
|
||||
)
|
||||
else:
|
||||
run_benchmark(
|
||||
args.seq_lens,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.method,
|
||||
)
|
||||
@@ -20,7 +20,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
deepgemm
|
||||
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
|
||||
GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
|
||||
GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
|
||||
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
@@ -120,11 +120,6 @@ if(DEEPGEMM_ARCHS)
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
|
||||
DESTINATION vllm/third_party/deep_gemm/mega
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
# Generate envs.py (normally generated by DeepGEMM's setup.py build step)
|
||||
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
|
||||
"# Pre-installed environment variables\npersistent_envs = dict()\n")
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
|
||||
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
+25
-82
@@ -11,74 +11,29 @@
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
bool act_first>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
}
|
||||
const scalar_t& y) {
|
||||
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
bool act_first>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
}
|
||||
const packed_t& y) {
|
||||
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
|
||||
: packed_mul(x, PACKED_ACT_FN(y));
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
bool use_vec, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d, const float limit) {
|
||||
const int d) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
|
||||
x.elts[j], y.elts[j]);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
|
||||
// Launch activation and gating kernel.
|
||||
// Use ACT_FIRST (bool) indicating whether to apply the activation function
|
||||
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
// first.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
true);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false, false, 0.0f);
|
||||
false);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true, false, 0.0f);
|
||||
true);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
+8
-18
@@ -599,11 +599,6 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
|
||||
// Find batch index within a block
|
||||
__shared__ int batch_idx[BLOCK_Y_SIZE];
|
||||
if (threadIdx.x == 0) {
|
||||
batch_idx[threadIdx.y] = -1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
|
||||
iter++) {
|
||||
int tid = iter * blockDim.x + threadIdx.x;
|
||||
@@ -616,18 +611,16 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
#ifndef USE_ROCM
|
||||
__syncwarp();
|
||||
#endif
|
||||
|
||||
// num_tokens may be an allocation upper bound when Python avoids a D2H sync.
|
||||
// Only tokens covered by the exact device-side cu_seq_lens are valid to
|
||||
// gather.
|
||||
const int batch = batch_idx[threadIdx.y];
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens) {
|
||||
return;
|
||||
}
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
|
||||
const int block_idx =
|
||||
block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
|
||||
const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
|
||||
inbatch_seq_idx / cache_block_size];
|
||||
const int64_t src_block_offset = block_idx * block_stride;
|
||||
const int64_t cache_inblock_offset =
|
||||
(inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
|
||||
@@ -1497,9 +1490,6 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
|
||||
TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
|
||||
ql_nope.scalar_type() == at::ScalarType::BFloat16,
|
||||
"ql_nope must be float16 or bfloat16 dtype");
|
||||
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
@@ -1511,7 +1501,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
|
||||
q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
|
||||
q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
|
||||
|
||||
@@ -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__
|
||||
|
||||
|
||||
@@ -1152,11 +1152,7 @@ class AttentionMainLoop {
|
||||
bool use_sink) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
|
||||
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
|
||||
|
||||
@@ -1205,17 +1201,8 @@ class AttentionMainLoop {
|
||||
vec = vec - max_vec;
|
||||
|
||||
// compute exp
|
||||
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
prob_buffer_vec_t output_vec(vec);
|
||||
output_vec.save(curr_prob_buffer_iter);
|
||||
#else
|
||||
@@ -1271,11 +1258,7 @@ class AttentionMainLoop {
|
||||
int32_t kv_tile_token_num, float softcap_scale) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
float inv_softcap_scale = 1.0 / softcap_scale;
|
||||
vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
|
||||
vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
|
||||
@@ -1289,15 +1272,8 @@ class AttentionMainLoop {
|
||||
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
|
||||
vec = vec * inv_softcap_scale_vec;
|
||||
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
|
||||
vec = (vec - inv_vec) / (vec + inv_vec);
|
||||
#else
|
||||
|
||||
@@ -15,12 +15,16 @@
|
||||
#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__)
|
||||
#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
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
@@ -482,18 +486,9 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
}
|
||||
|
||||
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);
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, 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 =
|
||||
@@ -711,18 +706,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
}
|
||||
|
||||
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);
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, 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 =
|
||||
|
||||
@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(
|
||||
|
||||
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
std::optional<torch::Tensor> key, int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, bool is_neox,
|
||||
int64_t rope_dim_offset, bool inverse) {
|
||||
TORCH_CHECK(rope_dim_offset == 0,
|
||||
"rope_dim_offset != 0 is not supported on CPU");
|
||||
TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
|
||||
|
||||
torch::Tensor& cos_sin_cache, bool is_neox) {
|
||||
int num_tokens = positions.numel();
|
||||
int rot_dim = cos_sin_cache.size(1);
|
||||
int num_heads = query.size(-1) / head_size;
|
||||
|
||||
@@ -263,8 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"rotary_embedding(Tensor positions, Tensor! query,"
|
||||
" Tensor!? key, int head_size,"
|
||||
" Tensor cos_sin_cache, bool is_neox, int "
|
||||
"rope_dim_offset=0, bool inverse=False) -> ()");
|
||||
" Tensor cos_sin_cache, bool is_neox) -> ()");
|
||||
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
|
||||
|
||||
// Quantization
|
||||
|
||||
@@ -1,477 +0,0 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* Horizontally-fused DeepseekV4-MLA kernel:
|
||||
* - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
|
||||
* - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
|
||||
* cache insert
|
||||
*
|
||||
* Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
|
||||
* TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
|
||||
* dispatch choosing Q vs KV work per warp. The per-warp RMSNorm/RoPE
|
||||
* skeleton is adapted from vllm-deepseek_v4's existing
|
||||
* `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
|
||||
*
|
||||
* Assumptions (hard-coded for DeepseekV4 attention):
|
||||
* HEAD_DIM = 512
|
||||
* ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
|
||||
* NOPE_DIM = 448
|
||||
* QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
|
||||
* FP8_MAX = 448.0f
|
||||
* is_neox=false (GPT-J interleaved pairs)
|
||||
* cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
|
||||
* second along last dim; each half is rope_dim/2 = 32 values)
|
||||
*
|
||||
* Cache layout per paged-cache block (block_size tokens):
|
||||
* [0, bs*576): token data, 448 fp8 + 128 bf16 each
|
||||
* [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/cuda.h>
|
||||
|
||||
#include "cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
#include "type_convert.cuh"
|
||||
|
||||
#ifndef FINAL_MASK
|
||||
#define FINAL_MASK 0xffffffffu
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace deepseek_v4_fused_ops {
|
||||
|
||||
namespace {
|
||||
inline int getSMVersion() {
|
||||
auto* props = at::cuda::getCurrentDeviceProperties();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Constants
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
constexpr int kHeadDim = 512;
|
||||
constexpr int kRopeDim = 64;
|
||||
constexpr int kNopeDim = kHeadDim - kRopeDim; // 448
|
||||
constexpr int kQuantBlock = 64;
|
||||
constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
|
||||
constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
|
||||
constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
|
||||
constexpr float kFp8Max = 448.0f;
|
||||
|
||||
// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
|
||||
constexpr int kNumLanes = 32;
|
||||
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Small inline helpers
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
__device__ __forceinline__ float warp4MaxAbs(float val) {
|
||||
// Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
|
||||
float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
|
||||
val = fmaxf(val, peer);
|
||||
peer = __shfl_xor_sync(FINAL_MASK, val, 2);
|
||||
val = fmaxf(val, peer);
|
||||
return val;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float warpSum(float val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1) {
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
|
||||
// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
|
||||
// warp handles one (token, head_slot) pair. head_slot < num_heads_q →
|
||||
// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV
|
||||
// branch (RoPE + UE8M0 quant + insert)
|
||||
//
|
||||
// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
|
||||
// The Q branch covers all `num_tokens_full` rows (downstream attention uses
|
||||
// them). The KV branch only inserts the first `num_tokens_insert` tokens
|
||||
// (= slot_mapping length) into the paged cache.
|
||||
//
|
||||
template <typename scalar_t_in>
|
||||
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
|
||||
scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
|
||||
scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
|
||||
uint8_t* __restrict__ k_cache, // [num_blocks, block_stride]
|
||||
int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
|
||||
int64_t const* __restrict__ position_ids, // [N] i64
|
||||
float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
|
||||
float const eps,
|
||||
int const num_tokens_full, // = q.size(0) = kv.size(0)
|
||||
int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
|
||||
int const num_heads_q, // H
|
||||
int const cache_block_size, // tokens per paged-cache block
|
||||
int const kv_block_stride) { // bytes per paged-cache block
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
// BF16 _typeConvert specialization is unavailable on pre-Ampere. The
|
||||
// DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
|
||||
// no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
|
||||
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
using Converter = vllm::_typeConvert<scalar_t_in>;
|
||||
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const warpId = threadIdx.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
|
||||
|
||||
int const total_slots_per_token = num_heads_q + 1;
|
||||
int const tokenIdx = globalWarpIdx / total_slots_per_token;
|
||||
int const slotIdx = globalWarpIdx % total_slots_per_token;
|
||||
if (tokenIdx >= num_tokens_full) return;
|
||||
|
||||
bool const isKV = (slotIdx == num_heads_q);
|
||||
// KV branch: skip DP-padded tokens (no slot reserved for them).
|
||||
if (isKV && tokenIdx >= num_tokens_insert) return;
|
||||
|
||||
// PDL: wait for predecessor kernel (upstream q/kv producer) to signal
|
||||
// before touching any global memory. No-op when PDL is not enabled on
|
||||
// the launch. The CUDA runtime wrapper emits the griddepcontrol.wait
|
||||
// PTX with the required memory clobber internally.
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// Dim range this lane owns within the 512-wide head.
|
||||
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
|
||||
|
||||
// ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
|
||||
float elements[kElemsPerLane];
|
||||
float sumOfSquares = 0.0f;
|
||||
|
||||
scalar_t_in const* src_ptr;
|
||||
if (isKV) {
|
||||
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
|
||||
} else {
|
||||
int64_t const q_row_offset =
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
src_ptr = q_inout + q_row_offset;
|
||||
}
|
||||
|
||||
// Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
|
||||
// type and bitcast to scalar_t_in packed pairs for conversion.
|
||||
uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
|
||||
uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
|
||||
|
||||
{
|
||||
typename Converter::packed_hip_type const* p0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
|
||||
typename Converter::packed_hip_type const* p1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
|
||||
// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p0[i]);
|
||||
elements[2 * i] = f2.x;
|
||||
elements[2 * i + 1] = f2.y;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p1[i]);
|
||||
elements[8 + 2 * i] = f2.x;
|
||||
elements[8 + 2 * i + 1] = f2.y;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
|
||||
// Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
|
||||
// The downstream bf16 round only happens at the final store.
|
||||
if (!isKV) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
sumOfSquares += elements[i] * elements[i];
|
||||
}
|
||||
sumOfSquares = warpSum<float>(sumOfSquares);
|
||||
float const rms_rcp =
|
||||
rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = elements[i] * rms_rcp;
|
||||
}
|
||||
}
|
||||
|
||||
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
|
||||
// All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
|
||||
bool const is_rope_lane = dim_base >= kNopeDim;
|
||||
if (is_rope_lane) {
|
||||
int64_t const pos = position_ids[tokenIdx];
|
||||
constexpr int kHalfRope = kRopeDim / 2; // 32
|
||||
float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
|
||||
float const* sin_ptr = cos_ptr + kHalfRope;
|
||||
|
||||
int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16
|
||||
#pragma unroll
|
||||
for (int p = 0; p < kElemsPerLane / 2; p++) {
|
||||
int const pair_dim = rope_local_base + 2 * p;
|
||||
int const half_idx = pair_dim / 2;
|
||||
float const cos_v = VLLM_LDG(cos_ptr + half_idx);
|
||||
float const sin_v = VLLM_LDG(sin_ptr + half_idx);
|
||||
float const x_even = elements[2 * p];
|
||||
float const x_odd = elements[2 * p + 1];
|
||||
elements[2 * p] = x_even * cos_v - x_odd * sin_v;
|
||||
elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Q branch: cast to bf16 and store back in place.
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
if (!isKV) {
|
||||
uint4 out0, out1;
|
||||
typename Converter::packed_hip_type* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
typename Converter::packed_hip_type* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst =
|
||||
q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// KV branch.
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
int64_t const slot_id = slot_mapping[tokenIdx];
|
||||
if (slot_id < 0) {
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t const block_idx = slot_id / cache_block_size;
|
||||
int64_t const pos_in_block = slot_id % cache_block_size;
|
||||
uint8_t* block_base =
|
||||
k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
|
||||
uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
|
||||
uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
|
||||
uint8_t* token_scale_ptr =
|
||||
block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
|
||||
pos_in_block * kScaleBytesPerToken;
|
||||
|
||||
// Round K to bf16 first, matching the unfused reference path where K is
|
||||
// materialized as bf16 before K quantization. absmax, clamp, and FP8
|
||||
// quant below all run on these bf16-rounded values.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = Converter::convert(Converter::convert(elements[i]));
|
||||
}
|
||||
|
||||
// Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
|
||||
// shuffle requires full participation). RoPE lanes contribute garbage,
|
||||
// but their values are gated out below via `!is_rope_lane`.
|
||||
float local_absmax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
|
||||
}
|
||||
float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
|
||||
float const exponent = ceilf(log2f(absmax / kFp8Max));
|
||||
float const inv_scale = exp2f(-exponent);
|
||||
|
||||
if (!is_rope_lane) {
|
||||
// ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
|
||||
uint8_t out_bytes[kElemsPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float scaled = elements[i] * inv_scale;
|
||||
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
|
||||
__nv_fp8_storage_t s =
|
||||
__nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
|
||||
out_bytes[i] = static_cast<uint8_t>(s);
|
||||
}
|
||||
// One 16-byte STG per lane.
|
||||
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
|
||||
*reinterpret_cast<uint4 const*>(out_bytes);
|
||||
|
||||
// Lane (4k) of each 4-lane group writes the scale byte for block k<7.
|
||||
if ((laneId & 3) == 0) {
|
||||
int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes
|
||||
float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
|
||||
token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
|
||||
}
|
||||
// Lane 0 also writes the padding byte at index 7.
|
||||
if (laneId == 0) {
|
||||
token_scale_ptr[kNumQuantBlocks] = 0; // pad
|
||||
}
|
||||
} else {
|
||||
// ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
|
||||
uint4 out0, out1;
|
||||
typename Converter::packed_hip_type* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
typename Converter::packed_hip_type* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
int const rope_local_base = dim_base - kNopeDim; // in [0, 64)
|
||||
scalar_t_in* bf16_dst =
|
||||
reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
|
||||
*reinterpret_cast<uint4*>(bf16_dst) = out0;
|
||||
*reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
|
||||
}
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Launch wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
template <typename scalar_t_in>
|
||||
void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
|
||||
int64_t const* slot_mapping, int64_t const* position_ids,
|
||||
float const* cos_sin_cache, float const eps, int const num_tokens_full,
|
||||
int const num_tokens_insert, int const num_heads_q,
|
||||
int const cache_block_size, int const kv_block_stride,
|
||||
cudaStream_t stream) {
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
|
||||
int const grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
|
||||
// PDL: enable programmatic stream serialization whenever the hardware
|
||||
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
|
||||
// so leave numAttrs = 0 and launch as a regular kernel.
|
||||
static int const sm_version = getSMVersion();
|
||||
// Host-side guard: the device kernel body is compiled as a no-op for
|
||||
// bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
|
||||
// unavailable there. Refuse the launch loudly instead of silently
|
||||
// skipping the work.
|
||||
TORCH_CHECK(
|
||||
sm_version >= 80,
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
|
||||
"(Ampere or newer); got sm_",
|
||||
sm_version);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = dim3(grid);
|
||||
config.blockDim = dim3(kBlockSize);
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = attrs;
|
||||
config.numAttrs = (sm_version >= 90) ? 1 : 0;
|
||||
|
||||
cudaLaunchKernelEx(
|
||||
&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
|
||||
q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
|
||||
num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
|
||||
kv_block_stride);
|
||||
}
|
||||
|
||||
} // namespace deepseek_v4_fused_ops
|
||||
} // namespace vllm
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Torch op wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::Tensor& q, // [N, H, 512] bf16, in place
|
||||
torch::Tensor const& kv, // [N, 512] bf16 (read-only)
|
||||
torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
|
||||
torch::Tensor const& slot_mapping, // [N] int64
|
||||
torch::Tensor const& position_ids, // [N] int64
|
||||
torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
|
||||
double eps, int64_t cache_block_size) {
|
||||
TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
|
||||
TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
|
||||
TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
|
||||
TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
|
||||
"position_ids must be int64 CUDA");
|
||||
TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
|
||||
TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
|
||||
TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
|
||||
TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64]");
|
||||
TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
|
||||
"cos_sin_cache must be float32");
|
||||
|
||||
// With DP padding, slot_mapping can be shorter than q/kv/positions.
|
||||
// Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
|
||||
// KV quant+insert runs only on the first slot_mapping.size(0) rows.
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
int const cache_block_size_i = static_cast<int>(cache_block_size);
|
||||
int const kv_block_stride = static_cast<int>(k_cache.stride(0));
|
||||
|
||||
at::cuda::OptionalCUDAGuard device_guard(device_of(q));
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
|
||||
using qkv_scalar_t = scalar_t;
|
||||
vllm::deepseek_v4_fused_ops::
|
||||
launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
|
||||
reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
|
||||
reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
|
||||
reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
|
||||
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
|
||||
cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
cache_block_size_i, kv_block_stride, stream);
|
||||
});
|
||||
}
|
||||
@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float x = static_cast<float>(src1.val[j]);
|
||||
float w = static_cast<float>(src2.val[j]);
|
||||
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
|
||||
dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
|
||||
}
|
||||
v_out[i] = dst;
|
||||
}
|
||||
@@ -135,17 +134,10 @@ fused_add_rms_norm_kernel(
|
||||
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
|
||||
int id = blockIdx.x * vec_hidden_size + idx;
|
||||
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
|
||||
_f16Vec<scalar_t, width> res = residual_v[id];
|
||||
_f16Vec<scalar_t, width> w = weight_v[idx];
|
||||
_f16Vec<scalar_t, width> out;
|
||||
using Converter = _typeConvert<scalar_t>;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < width; ++j) {
|
||||
float x = Converter::convert(res.data[j]);
|
||||
float wf = Converter::convert(w.data[j]);
|
||||
out.data[j] = Converter::convert(x * s_variance * wf);
|
||||
}
|
||||
input_v[strided_id] = out;
|
||||
_f16Vec<scalar_t, width> temp = residual_v[id];
|
||||
temp *= s_variance;
|
||||
temp *= weight_v[idx];
|
||||
input_v[strided_id] = temp;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -182,8 +174,8 @@ fused_add_rms_norm_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
|
||||
float x = (float)residual[blockIdx.x * hidden_size + idx];
|
||||
float w = (float)weight[idx];
|
||||
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
|
||||
input[blockIdx.x * input_stride + idx] =
|
||||
((scalar_t)(x * s_variance)) * weight[idx];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float x = static_cast<float>(src1.val[j]);
|
||||
float w = static_cast<float>(src2.val[j]);
|
||||
// Round normalized result through scalar_t to match the precision of the
|
||||
// unfused composite (rms_norm writes scalar_t, then
|
||||
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
|
||||
// Without this round, the fused path is strictly more accurate and
|
||||
// disagrees with the composite at exact E4M3 quantization tie boundaries.
|
||||
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
|
||||
float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
|
||||
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
|
||||
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
|
||||
scale_inv);
|
||||
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
|
||||
int id = blockIdx.x * vec_hidden_size + idx;
|
||||
_f16Vec<scalar_t, width> res = residual_v[id];
|
||||
_f16Vec<scalar_t, width> w = weight_v[idx];
|
||||
using Converter = _typeConvert<scalar_t>;
|
||||
using HipT = typename Converter::hip_type;
|
||||
_f16Vec<scalar_t, width> temp = residual_v[id];
|
||||
temp *= s_variance;
|
||||
temp *= weight_v[idx];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < width; ++i) {
|
||||
float x = Converter::convert(res.data[i]);
|
||||
float wf = Converter::convert(w.data[i]);
|
||||
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
|
||||
// to match the unfused composite path at FP8 boundaries. We use the
|
||||
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
|
||||
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
|
||||
HipT out_norm_h = Converter::convert(x * s_variance * wf);
|
||||
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
|
||||
Converter::convert(out_norm_h), scale_inv);
|
||||
out[id * width + i] =
|
||||
scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
|
||||
float x = (float)residual[blockIdx.x * hidden_size + idx];
|
||||
float w = (float)weight[idx];
|
||||
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
|
||||
// to match the unfused composite path at FP8 boundaries.
|
||||
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
|
||||
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
|
||||
static_cast<float>(out_norm), scale_inv);
|
||||
float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
|
||||
out[blockIdx.x * hidden_size + idx] =
|
||||
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -277,9 +277,7 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
|
||||
if (blockRepeat > 1) {
|
||||
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
|
||||
// The shared-memory vectorized offset load only handles full 4-expert
|
||||
// chunks. Use the scalar specialization for the remainder cases.
|
||||
if (n_experts >= 4 && n_experts % 4 == 0) {
|
||||
if (n_experts >= 4) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
@@ -301,9 +299,7 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
n_experts);
|
||||
}
|
||||
} else {
|
||||
// The low-latency vectorized expert lookup only handles full 16-expert
|
||||
// chunks. Fall back to the scalar lookup path for the remainder cases.
|
||||
if (n_experts >= 16 && n_experts % 16 == 0) {
|
||||
if (n_experts >= 16) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
|
||||
@@ -12,15 +12,6 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
|
||||
torch::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::Tensor> bias);
|
||||
|
||||
void topk_softplus_sqrt(torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices,
|
||||
torch::Tensor& token_expert_indices,
|
||||
torch::Tensor& gating_output, bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid);
|
||||
|
||||
void moe_sum(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
|
||||
|
||||
@@ -1,715 +0,0 @@
|
||||
/*
|
||||
* Adapted from
|
||||
* https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
|
||||
* Copyright (c) 2024, The vLLM team.
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include <type_traits>
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "../cuda_compat.h"
|
||||
#include "../cub_helpers.h"
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#else
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
typedef __hip_bfloat16 __nv_bfloat16;
|
||||
typedef __hip_bfloat162 __nv_bfloat162;
|
||||
#endif
|
||||
|
||||
#define MAX(a, b) ((a) > (b) ? (a) : (b))
|
||||
#define MIN(a, b) ((a) < (b) ? (a) : (b))
|
||||
|
||||
namespace vllm {
|
||||
namespace moe {
|
||||
|
||||
/// Aligned array type
|
||||
template <typename T,
|
||||
/// Number of elements in the array
|
||||
int N,
|
||||
/// Alignment requirement in bytes
|
||||
int Alignment = sizeof(T) * N>
|
||||
struct alignas(Alignment) AlignedArray {
|
||||
T data[N];
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float toFloat(T value) {
|
||||
if constexpr (std::is_same_v<T, float>) {
|
||||
return value;
|
||||
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
|
||||
return __bfloat162float(value);
|
||||
} else if constexpr (std::is_same_v<T, __half>) {
|
||||
return __half2float(value);
|
||||
}
|
||||
}
|
||||
|
||||
#define FINAL_MASK 0xffffffff
|
||||
template <typename T>
|
||||
__inline__ __device__ T warpReduceSum(T val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
return val;
|
||||
}
|
||||
|
||||
// ====================== TopK softplus_sqrt things
|
||||
// ===============================
|
||||
|
||||
/*
|
||||
A Top-K gating softplus_sqrt written to exploit when the number of experts in
|
||||
the MoE layers are a small power of 2. This allows us to cleanly share the
|
||||
rows among the threads in a single warp and eliminate communication between
|
||||
warps (so no need to use shared mem).
|
||||
|
||||
It fuses the sigmoid, max and argmax into a single kernel.
|
||||
|
||||
Limitations:
|
||||
1) This implementation is optimized for when the number of experts is a small
|
||||
power of 2. Additionally it also supports when number of experts is multiple
|
||||
of 64 which is still faster than the computing sigmoid and topK separately
|
||||
(only tested on CUDA yet). 2) This implementation assumes k is small, but will
|
||||
work for any k.
|
||||
*/
|
||||
|
||||
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
|
||||
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
|
||||
typename InputType = float>
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGatingSoftplusSqrt(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const IndType* input_ids, const IndType* tid2eid) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
"InputType must be float, __nv_bfloat16, or __half");
|
||||
|
||||
// We begin by enforcing compile time assertions and setting up compile time
|
||||
// constants.
|
||||
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
|
||||
"BYTES_PER_LDG must be power of 2");
|
||||
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
|
||||
|
||||
// Number of bytes each thread pulls in per load
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
|
||||
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
|
||||
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
|
||||
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
|
||||
|
||||
if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>) {
|
||||
static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
|
||||
"ELTS_PER_LDG must be 1 or even for 16-bit conversion");
|
||||
}
|
||||
|
||||
// Restrictions based on previous section.
|
||||
static_assert(
|
||||
VPT % ELTS_PER_LDG == 0,
|
||||
"The elements per thread must be a multiple of the elements per ldg");
|
||||
static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
|
||||
"The threads per row must cleanly divide the threads per warp");
|
||||
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
|
||||
"THREADS_PER_ROW must be power of 2");
|
||||
static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
|
||||
"THREADS_PER_ROW can be at most warp size");
|
||||
|
||||
// We have NUM_EXPERTS elements per row. We specialize for small #experts
|
||||
static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
|
||||
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
|
||||
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
|
||||
|
||||
// Restrictions for previous section.
|
||||
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
|
||||
"The elts per row must cleanly divide the total elt per warp");
|
||||
|
||||
// ===================== From this point, we finally start computing run-time
|
||||
// variables. ========================
|
||||
|
||||
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a
|
||||
// block contains WARPS_PER_CTA warps. This, each block processes a chunk of
|
||||
// rows. We start by computing the start row for each block.
|
||||
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
|
||||
|
||||
// Now, using the base row per thread block, we compute the base row per warp.
|
||||
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
|
||||
|
||||
// The threads in a warp are split into sub-groups that will work on a row.
|
||||
// We compute row offset for each thread sub-group
|
||||
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
|
||||
const int thread_row = warp_base_row + thread_row_in_warp;
|
||||
|
||||
// Threads with indices out of bounds should early exit here.
|
||||
if (thread_row >= num_rows) {
|
||||
return;
|
||||
}
|
||||
const bool row_is_active = finished ? !finished[thread_row] : true;
|
||||
|
||||
// We finally start setting up the read pointers for each thread. First, each
|
||||
// thread jumps to the start of the row it will read.
|
||||
const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
|
||||
|
||||
// Now, we compute the group each thread belong to in order to determine the
|
||||
// first column to start loads.
|
||||
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
|
||||
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
|
||||
const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
|
||||
|
||||
// Finally, we pull in the data from global mem
|
||||
float row_chunk[VPT];
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
// to float
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
using VecType = AlignedArray<float, ELTS_PER_LDG>;
|
||||
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
|
||||
if constexpr (ELTS_PER_LDG >= 2) {
|
||||
using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
|
||||
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
|
||||
#pragma unroll
|
||||
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
|
||||
row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
|
||||
*reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
|
||||
}
|
||||
}
|
||||
} else { // ELTS_PER_LDG == 1
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
const __nv_bfloat16* scalar_ptr =
|
||||
thread_read_ptr + ii * THREADS_PER_ROW;
|
||||
row_chunk[ii] = __bfloat162float(*scalar_ptr);
|
||||
}
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, __half>) {
|
||||
if constexpr (ELTS_PER_LDG >= 2) {
|
||||
using VecType = AlignedArray<__half, ELTS_PER_LDG>;
|
||||
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
|
||||
#pragma unroll
|
||||
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
|
||||
row_chunk_f2[base_idx_f2 + jj] = __half22float2(
|
||||
*reinterpret_cast<const __half2*>(vec.data + jj * 2));
|
||||
}
|
||||
}
|
||||
} else { // ELTS_PER_LDG == 1
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
|
||||
row_chunk[ii] = __half2float(*scalar_ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
constexpr float threshold = 20.0f;
|
||||
constexpr float beta = 1.0f;
|
||||
|
||||
// Hash MoE path: indices are predetermined from lookup table
|
||||
if constexpr (USE_HASH) {
|
||||
const IndType token_id = input_ids[thread_row];
|
||||
const IndType* expert_indices_for_token = tid2eid + token_id * k;
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
row_chunk[ii] = sqrtf(val);
|
||||
}
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = expert;
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Compute per-thread scale (using warp reduction when renormalizing).
|
||||
if (renormalize) {
|
||||
selected_sum = warpReduceSum(selected_sum);
|
||||
}
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
scale /= denom;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
output[idx] = row_chunk[ii] * scale;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
|
||||
// Original TopK path: find top-k experts by score
|
||||
// Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
|
||||
// the topk elements in each row, along with the max index.
|
||||
int start_col = first_elt_read_by_thread;
|
||||
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
|
||||
|
||||
float selected_sum = 0.f;
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
// First, each thread does the local argmax
|
||||
float max_val = row_chunk[0];
|
||||
int expert = start_col;
|
||||
#pragma unroll
|
||||
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
|
||||
++ldg, col += COLS_PER_GROUP_LDG) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
|
||||
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
|
||||
|
||||
// No check on the experts here since columns with the smallest index
|
||||
// are processed first and only updated if > (not >=)
|
||||
if (val > max_val) {
|
||||
max_val = val;
|
||||
expert = col + ii;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
|
||||
// reach consensus about the max. This will be useful for K > 1 so that the
|
||||
// threads can agree on "who" had the max value. That thread can then blank out
|
||||
// their max with -inf and the warp can run more iterations...
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
float other_max =
|
||||
VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
|
||||
int other_expert =
|
||||
VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
|
||||
|
||||
// We want lower indices to "win" in every thread so we break ties this
|
||||
// way
|
||||
if (other_max > max_val ||
|
||||
(other_max == max_val && other_expert < expert)) {
|
||||
max_val = other_max;
|
||||
expert = other_expert;
|
||||
}
|
||||
}
|
||||
|
||||
// Write the max for this k iteration to global memory.
|
||||
if (thread_group_idx == 0) {
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert =
|
||||
expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
|
||||
// The lead thread from each sub-group will write out the final results to
|
||||
// global memory. (This will be a single) thread per row of the
|
||||
// input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
if (correction_bias != nullptr) {
|
||||
max_val -= correction_bias[expert];
|
||||
}
|
||||
output[idx] = max_val;
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
}
|
||||
}
|
||||
|
||||
// Finally, we clear the value in the thread with the current max if there
|
||||
// is another iteration to run.
|
||||
if (k_idx + 1 < k) {
|
||||
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
|
||||
const int thread_to_clear_in_group =
|
||||
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
|
||||
|
||||
// Only the thread in the group which produced the max will reset the
|
||||
// "winning" value to -inf.
|
||||
if (thread_group_idx == thread_to_clear_in_group) {
|
||||
const int offset_for_expert = expert % ELTS_PER_LDG;
|
||||
// Safe to set to any negative value since row_chunk values must be
|
||||
// between 0 and 1.
|
||||
row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
|
||||
-10000.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply renormalization and routed scaling factor to final weights.
|
||||
if (thread_group_idx == 0) {
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
scale /= denom;
|
||||
}
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = output[idx] * scale;
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
namespace detail {
|
||||
// Constructs some constants needed to partition the work across threads at
|
||||
// compile time.
|
||||
template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
|
||||
typename InputType>
|
||||
struct TopkConstants {
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
|
||||
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
|
||||
EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
|
||||
"");
|
||||
static constexpr int VECs_PER_THREAD =
|
||||
MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
|
||||
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
|
||||
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
|
||||
static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
#define DISPATCH_HASH(use_hash, USE_HASH, ...) \
|
||||
if (use_hash) { \
|
||||
const bool USE_HASH = true; \
|
||||
static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
const bool USE_HASH = false; \
|
||||
static_assert(USE_HASH == false, \
|
||||
"USE_HASH must be compile-time constant"); \
|
||||
__VA_ARGS__ \
|
||||
}
|
||||
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
|
||||
int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
|
||||
void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
IndType* indices, int* source_row, const int num_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
static constexpr int VPT = Constants::VPT;
|
||||
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
|
||||
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
|
||||
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
DISPATCH_HASH(use_hash, USE_HASH, {
|
||||
auto* kernel =
|
||||
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
|
||||
WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
|
||||
#ifndef USE_ROCM
|
||||
cudaLaunchConfig_t config = {};
|
||||
config.gridDim = num_blocks;
|
||||
config.blockDim = block_dim;
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.numAttrs = 1;
|
||||
config.attrs = attrs;
|
||||
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
|
||||
indices, source_row, k, start_expert, end_expert,
|
||||
renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
#else
|
||||
kernel<<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert,
|
||||
end_expert, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
#endif
|
||||
})
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
static_assert(WARP_SIZE == 32, \
|
||||
"Unsupported warp size. Only 32 is supported for CUDA"); \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
|
||||
stream);
|
||||
#else
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename IndType, typename InputType>
|
||||
void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
|
||||
int* token_expert_indices, const int num_tokens, const int num_experts,
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash, const IndType* input_ids,
|
||||
const IndType* tid2eid, cudaStream_t stream) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
|
||||
// elements can be loaded by a warp
|
||||
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
|
||||
(std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>)
|
||||
? 4
|
||||
: 8;
|
||||
#endif
|
||||
switch (num_experts) {
|
||||
case 1:
|
||||
LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 2:
|
||||
LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 4:
|
||||
LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 8:
|
||||
LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 16:
|
||||
LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 32:
|
||||
LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 64:
|
||||
LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 128:
|
||||
LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 256:
|
||||
LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 512:
|
||||
LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
|
||||
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
|
||||
// num_experts, alternatively we can test 4 bytes loading and enable it in
|
||||
// future.
|
||||
#ifndef USE_ROCM
|
||||
case 192:
|
||||
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 320:
|
||||
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 384:
|
||||
LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 448:
|
||||
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 576:
|
||||
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
#endif
|
||||
default: {
|
||||
TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe
|
||||
} // namespace vllm
|
||||
|
||||
template <typename ComputeType>
|
||||
void dispatch_topk_softplus_sqrt_launch(
|
||||
const ComputeType* gating_output, torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().data_ptr<float>();
|
||||
}
|
||||
bool use_hash = false;
|
||||
if (tid2eid.has_value()) {
|
||||
TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
|
||||
use_hash = true;
|
||||
}
|
||||
if (topk_indices.scalar_type() == at::ScalarType::Int) {
|
||||
const int* input_ids_ptr = nullptr;
|
||||
const int* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<int>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<int>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
} else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
|
||||
const uint32_t* input_ids_ptr = nullptr;
|
||||
const uint32_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
|
||||
}
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
} else {
|
||||
TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
|
||||
|
||||
const int64_t* input_ids_ptr = nullptr;
|
||||
const int64_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<int64_t>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
}
|
||||
}
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::Tensor& topk_weights, // [num_tokens, topk]
|
||||
torch::Tensor& topk_indices, // [num_tokens, topk]
|
||||
torch::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid) {
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
if (gating_output.scalar_type() == at::ScalarType::Float) {
|
||||
dispatch_topk_softplus_sqrt_launch<float>(
|
||||
gating_output.data_ptr<float>(), topk_weights, topk_indices,
|
||||
token_expert_indices, num_tokens, num_experts, topk, renormalize,
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
|
||||
} else if (gating_output.scalar_type() == at::ScalarType::Half) {
|
||||
dispatch_topk_softplus_sqrt_launch<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
} else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
|
||||
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(
|
||||
gating_output.data_ptr<at::BFloat16>()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
}
|
||||
}
|
||||
@@ -16,14 +16,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"bias) -> ()");
|
||||
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
|
||||
#endif
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts.
|
||||
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
|
||||
|
||||
+1
-9
@@ -100,11 +100,6 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
|
||||
bool is_neox, torch::Tensor& position_ids,
|
||||
int64_t forced_token_heads_per_warp);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
|
||||
torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
|
||||
torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
|
||||
|
||||
void apply_repetition_penalties_(torch::Tensor& logits,
|
||||
const torch::Tensor& prompt_mask,
|
||||
const torch::Tensor& output_mask,
|
||||
@@ -158,13 +153,10 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
|
||||
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
std::optional<torch::Tensor> key, int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, bool is_neox,
|
||||
int64_t rope_dim_offset, bool inverse);
|
||||
torch::Tensor& cos_sin_cache, bool is_neox);
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
|
||||
+16
-17
@@ -18,6 +18,7 @@ namespace persistent {
|
||||
// Constants
|
||||
// ============================================================================
|
||||
|
||||
constexpr int TopK = 2048;
|
||||
constexpr int kThreadsPerBlock = 1024;
|
||||
constexpr int RADIX = 256;
|
||||
|
||||
@@ -127,12 +128,11 @@ struct RadixRowState {
|
||||
|
||||
struct PersistentTopKParams {
|
||||
const float* __restrict__ input; // [num_rows, stride]
|
||||
int32_t* __restrict__ output; // [num_rows, top_k]
|
||||
int32_t* __restrict__ output; // [num_rows, TopK]
|
||||
int32_t* __restrict__ lengths; // [num_rows]
|
||||
RadixRowState* row_states; // large path: per-group state
|
||||
uint32_t num_rows;
|
||||
uint32_t stride;
|
||||
uint32_t top_k; // actual k value for output stride
|
||||
uint32_t chunk_size; // large path: elements per CTA
|
||||
uint32_t ctas_per_group; // 1=medium, >1=large
|
||||
uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit)
|
||||
@@ -154,7 +154,6 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
|
||||
return key >> 5;
|
||||
}
|
||||
|
||||
template <int TopK>
|
||||
__device__ __noinline__ void histogram_2048_topk(
|
||||
const float* __restrict__ logits, int32_t* __restrict__ output_indices,
|
||||
int32_t seq_len) {
|
||||
@@ -419,7 +418,6 @@ __device__ __noinline__ void histogram_2048_topk(
|
||||
// by: DarkSharpness
|
||||
// which at the same time is an optimized topk kernel copied from tilelang
|
||||
// kernel
|
||||
template <int TopK>
|
||||
__device__ __noinline__ void histogram_256_topk(
|
||||
const float* __restrict__ logits, int* __restrict__ output_indices,
|
||||
int logits_offset, int seq_len) {
|
||||
@@ -651,7 +649,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
|
||||
// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
|
||||
// ============================================================================
|
||||
|
||||
template <int TopK, uint32_t VEC_SIZE>
|
||||
template <uint32_t VEC_SIZE>
|
||||
__device__ void radix_topk(const float* __restrict__ row_input,
|
||||
int32_t* __restrict__ row_output, uint32_t seq_len,
|
||||
uint32_t my_chunk_start, uint32_t chunk_size,
|
||||
@@ -859,7 +857,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
|
||||
// see filtered_topk.cuh)
|
||||
// ============================================================================
|
||||
|
||||
template <int TopK = 2048, uint32_t VEC_SIZE = 1>
|
||||
template <uint32_t VEC_SIZE = 1>
|
||||
__global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
persistent_topk_kernel(PersistentTopKParams params) {
|
||||
const uint32_t tx = threadIdx.x;
|
||||
@@ -917,7 +915,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
if (row_idx >= params.num_rows) break;
|
||||
|
||||
const uint32_t seq_len = params.lengths[row_idx];
|
||||
int32_t* row_output = params.output + row_idx * params.top_k;
|
||||
int32_t* row_output = params.output + row_idx * TopK;
|
||||
const float* row_input = params.input + row_idx * params.stride;
|
||||
|
||||
if (seq_len <= RADIX_THRESHOLD) {
|
||||
@@ -929,19 +927,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
|
||||
}
|
||||
} else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
|
||||
histogram_2048_topk<TopK>(row_input, row_output, seq_len);
|
||||
histogram_2048_topk(row_input, row_output, seq_len);
|
||||
} else {
|
||||
histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
|
||||
histogram_256_topk(row_input, row_output, 0, seq_len);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const uint32_t my_chunk_start = cta_in_group * chunk_size;
|
||||
radix_topk<TopK, VEC_SIZE>(
|
||||
row_input, row_output, seq_len, my_chunk_start, chunk_size,
|
||||
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
|
||||
cta_in_group, ctas_per_group, barrier_phase, iter, tx);
|
||||
radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
|
||||
chunk_size, local_histogram, suffix_sum,
|
||||
shared_scalars, shared_ordered, state, cta_in_group,
|
||||
ctas_per_group, barrier_phase, iter, tx);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1013,6 +1011,7 @@ struct FilteredTopKTraits<float> {
|
||||
}
|
||||
};
|
||||
|
||||
constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
|
||||
constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
|
||||
constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
|
||||
16 * 1024; // 16K indices per buffer
|
||||
@@ -1026,7 +1025,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
|
||||
* \tparam IdType Index type (int32_t)
|
||||
* \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
|
||||
*/
|
||||
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
|
||||
template <typename DType, typename IdType, int VEC_SIZE>
|
||||
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
|
||||
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
|
||||
IdType* __restrict__ output,
|
||||
@@ -1060,7 +1059,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
|
||||
alignas(128) __shared__ int s_counter;
|
||||
alignas(128) __shared__ int s_threshold_bin_id;
|
||||
alignas(128) __shared__ int s_num_input[2];
|
||||
alignas(128) __shared__ int s_indices[MAX_K];
|
||||
alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
|
||||
|
||||
auto& s_histogram = s_histogram_buf[0];
|
||||
|
||||
@@ -1281,7 +1280,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
|
||||
return static_cast<int>(g);
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType, uint32_t MAX_K = 2048>
|
||||
template <typename DType, typename IdType>
|
||||
cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
|
||||
IdType* lengths, uint32_t num_rows,
|
||||
uint32_t top_k_val, uint32_t max_len,
|
||||
@@ -1298,7 +1297,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
|
||||
|
||||
#define DISPATCH_VEC_SIZE(VS) \
|
||||
if (vec_size == VS) { \
|
||||
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
|
||||
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>; \
|
||||
FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
|
||||
FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
|
||||
|
||||
@@ -9,29 +9,28 @@ namespace vllm {
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
inline __device__ void apply_token_rotary_embedding(
|
||||
scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
|
||||
const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
|
||||
const bool inverse) {
|
||||
scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
|
||||
const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
|
||||
int x_index, y_index;
|
||||
float cos_f, sin_f;
|
||||
scalar_t cos, sin;
|
||||
if (IS_NEOX) {
|
||||
// GPT-NeoX style rotary embedding.
|
||||
x_index = rot_offset;
|
||||
y_index = embed_dim + rot_offset;
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index);
|
||||
cos = VLLM_LDG(cos_ptr + x_index);
|
||||
sin = VLLM_LDG(sin_ptr + x_index);
|
||||
} else {
|
||||
// GPT-J style rotary embedding.
|
||||
x_index = 2 * rot_offset;
|
||||
y_index = 2 * rot_offset + 1;
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index / 2);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index / 2);
|
||||
cos = VLLM_LDG(cos_ptr + x_index / 2);
|
||||
sin = VLLM_LDG(sin_ptr + x_index / 2);
|
||||
}
|
||||
if (inverse) {
|
||||
sin_f = -sin_f;
|
||||
}
|
||||
const float x_f = static_cast<float>(arr[x_index]);
|
||||
const float y_f = static_cast<float>(arr[y_index]);
|
||||
arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
|
||||
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
|
||||
|
||||
const scalar_t x = arr[x_index];
|
||||
const scalar_t y = arr[y_index];
|
||||
arr[x_index] = x * cos - y * sin;
|
||||
arr[y_index] = y * cos + x * sin;
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
@@ -43,23 +42,22 @@ inline __device__ void apply_rotary_embedding(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const float* cache_ptr, const int head_size, const int num_heads,
|
||||
const scalar_t* cache_ptr, const int head_size, const int num_heads,
|
||||
const int num_kv_heads, const int rot_dim, const int token_idx,
|
||||
const int64_t query_stride, const int64_t key_stride,
|
||||
const int64_t head_stride, const int64_t rope_dim_offset,
|
||||
const bool inverse) {
|
||||
const int64_t head_stride) {
|
||||
const int embed_dim = rot_dim / 2;
|
||||
const float* cos_ptr = cache_ptr;
|
||||
const float* sin_ptr = cache_ptr + embed_dim;
|
||||
const scalar_t* cos_ptr = cache_ptr;
|
||||
const scalar_t* sin_ptr = cache_ptr + embed_dim;
|
||||
|
||||
const int nq = num_heads * embed_dim;
|
||||
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
|
||||
const int head_idx = i / embed_dim;
|
||||
const int64_t token_head =
|
||||
token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
|
||||
token_idx * query_stride + head_idx * head_stride;
|
||||
const int rot_offset = i % embed_dim;
|
||||
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
||||
}
|
||||
|
||||
if (key != nullptr) {
|
||||
@@ -67,10 +65,10 @@ inline __device__ void apply_rotary_embedding(
|
||||
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
|
||||
const int head_idx = i / embed_dim;
|
||||
const int64_t token_head =
|
||||
token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
|
||||
token_idx * key_stride + head_idx * head_stride;
|
||||
const int rot_offset = i % embed_dim;
|
||||
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -86,18 +84,19 @@ __global__ void rotary_embedding_kernel(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
|
||||
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
|
||||
// 2]
|
||||
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
|
||||
const int64_t head_stride, const int num_heads, const int num_kv_heads,
|
||||
const int head_size, const int64_t rope_dim_offset, const bool inverse) {
|
||||
const int head_size) {
|
||||
// Each thread block is responsible for one token.
|
||||
const int token_idx = blockIdx.x;
|
||||
int64_t pos = positions[token_idx];
|
||||
const float* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
|
||||
apply_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
|
||||
token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
|
||||
inverse);
|
||||
token_idx, query_stride, key_stride, head_stride);
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
@@ -116,7 +115,7 @@ void rotary_embedding(
|
||||
// [num_tokens, num_heads, head_size]
|
||||
int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
|
||||
bool is_neox, int64_t rope_dim_offset, bool inverse) {
|
||||
bool is_neox) {
|
||||
// num_tokens = batch_size * seq_len
|
||||
int64_t num_tokens = positions.numel();
|
||||
int positions_ndim = positions.dim();
|
||||
@@ -155,8 +154,6 @@ void rotary_embedding(
|
||||
int seq_dim_idx = positions_ndim - 1;
|
||||
int64_t query_stride = query.stride(seq_dim_idx);
|
||||
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
|
||||
|
||||
TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
|
||||
// Determine head stride: for [*, heads, head_size] use stride of last dim;
|
||||
// for flat [*, heads*head_size], heads blocks are contiguous of size
|
||||
// head_size
|
||||
@@ -168,23 +165,20 @@ void rotary_embedding(
|
||||
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
|
||||
VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
|
||||
if (is_neox) {
|
||||
vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
|
||||
inverse);
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size);
|
||||
} else {
|
||||
vllm::rotary_embedding_kernel<scalar_t, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
|
||||
inverse);
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
|
||||
key_stride, head_stride, num_heads, num_kv_heads, head_size);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
+11
-2
@@ -40,6 +40,15 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
|
||||
#define __HIP__FP8MFMA__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
|
||||
defined(__gfx1150__) || defined(__gfx1151__))
|
||||
#define __HIP__GFX11__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(NDEBUG)
|
||||
#undef NDEBUG
|
||||
#include <assert.h>
|
||||
@@ -1620,7 +1629,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
#elif defined(__GFX11__)
|
||||
#elif defined(__HIP__GFX11__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
@@ -2379,7 +2388,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
|
||||
}
|
||||
|
||||
#elif defined(__GFX12__)
|
||||
#elif defined(__HIP__GFX12__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
|
||||
+23
-18
@@ -26,11 +26,16 @@
|
||||
#define __HIP__GFX9__
|
||||
#endif
|
||||
|
||||
// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
|
||||
#if defined(__GFX11__) || defined(__GFX12__)
|
||||
#if defined(__HIPCC__) && \
|
||||
(defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
|
||||
defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX1X__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
|
||||
#define __HIP__MI3XX__
|
||||
#endif
|
||||
@@ -1840,7 +1845,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
return out_c;
|
||||
}
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -1888,7 +1893,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -1926,7 +1931,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -1950,7 +1955,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -1988,7 +1993,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2004,7 +2009,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2022,7 +2027,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2034,9 +2039,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
const int _WvPrGrp, const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -2083,7 +2088,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -2123,7 +2128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -2147,7 +2152,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -2185,7 +2190,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2201,7 +2206,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __GFX12__
|
||||
#ifdef __HIP__GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2219,7 +2224,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2231,7 +2236,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
|
||||
void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
|
||||
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
|
||||
|
||||
+1
-7
@@ -258,13 +258,7 @@ __device__ bool processHistogramStep(
|
||||
auto processBins = [&](float logit, int idx) {
|
||||
if (isPartialMatch<patternShift>(logit, logitPattern)) {
|
||||
uint32_t binIdx = extractBinIdx<step>(logit);
|
||||
// Only write elements with binIdx < thresholdBinIdx when:
|
||||
// 1. This is step 0 and the threshold bin is small enough (no step 1)
|
||||
// 2. This is step >= 1 (where pattern matching filters correctly)
|
||||
// This prevents duplicates when step 0 and step 1 both run.
|
||||
bool shouldWriteDirectly =
|
||||
(step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
|
||||
if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
|
||||
if (binIdx < thresholdBinIdx) {
|
||||
// The element is part of the top-k selection
|
||||
int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
|
||||
|
||||
|
||||
+35
-59
@@ -10,17 +10,33 @@
|
||||
#include "persistent_topk.cuh"
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len) {
|
||||
#ifndef USE_ROCM
|
||||
template <int TopK>
|
||||
void launch_persistent_topk(const torch::Tensor& logits,
|
||||
const torch::Tensor& lengths, torch::Tensor& output,
|
||||
torch::Tensor& workspace, int64_t max_seq_len) {
|
||||
namespace P = vllm::persistent;
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
|
||||
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
|
||||
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.size(1);
|
||||
|
||||
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
namespace P = vllm::persistent;
|
||||
|
||||
TORCH_CHECK(k == P::TopK, "k must be 2048");
|
||||
TORCH_CHECK(k <= stride, "k out of range");
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
static int num_sms = 0;
|
||||
@@ -34,17 +50,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
}
|
||||
|
||||
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
|
||||
cudaError_t status =
|
||||
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
|
||||
cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
|
||||
TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK failed: ", cudaGetErrorString(status));
|
||||
} else {
|
||||
TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
|
||||
TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
|
||||
|
||||
// Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
|
||||
// large path. Empirically tuned.
|
||||
int effective_max_smem;
|
||||
if (num_rows <= 4) {
|
||||
effective_max_smem =
|
||||
@@ -84,7 +101,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
|
||||
int occupancy = 1;
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
&occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
if (occupancy < 1) occupancy = 1;
|
||||
|
||||
@@ -104,16 +121,15 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
params.lengths = lengths.data_ptr<int32_t>();
|
||||
params.num_rows = static_cast<uint32_t>(num_rows);
|
||||
params.stride = static_cast<uint32_t>(stride);
|
||||
params.top_k = static_cast<uint32_t>(TopK);
|
||||
params.chunk_size = chunk_size;
|
||||
params.row_states =
|
||||
reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
|
||||
params.ctas_per_group = ctas_per_group;
|
||||
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
|
||||
|
||||
#define LAUNCH_PERSISTENT(TOPK_VAL, VS) \
|
||||
#define LAUNCH_PERSISTENT(VS) \
|
||||
do { \
|
||||
auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
|
||||
auto kernel = &P::persistent_topk_kernel<VS>; \
|
||||
cudaError_t err = cudaFuncSetAttribute( \
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
|
||||
TORCH_CHECK(err == cudaSuccess, \
|
||||
@@ -122,11 +138,11 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
} while (0)
|
||||
|
||||
if (vec_size == 4) {
|
||||
LAUNCH_PERSISTENT(TopK, 4);
|
||||
LAUNCH_PERSISTENT(4);
|
||||
} else if (vec_size == 2) {
|
||||
LAUNCH_PERSISTENT(TopK, 2);
|
||||
LAUNCH_PERSISTENT(2);
|
||||
} else {
|
||||
LAUNCH_PERSISTENT(TopK, 1);
|
||||
LAUNCH_PERSISTENT(1);
|
||||
}
|
||||
#undef LAUNCH_PERSISTENT
|
||||
}
|
||||
@@ -134,46 +150,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
cudaError_t err = cudaGetLastError();
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"persistent_topk failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
#endif
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len) {
|
||||
#ifndef USE_ROCM
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
|
||||
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
|
||||
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.size(1);
|
||||
|
||||
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
|
||||
"persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
|
||||
|
||||
if (k == 512) {
|
||||
launch_persistent_topk<512>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
} else if (k == 1024) {
|
||||
launch_persistent_topk<1024>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
} else {
|
||||
launch_persistent_topk<2048>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
}
|
||||
#else
|
||||
TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
|
||||
#endif
|
||||
|
||||
+1
-21
@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
|
||||
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
@@ -183,19 +177,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"int forced_token_heads_per_warp=-1) -> ()");
|
||||
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch.
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
|
||||
"Tensor! q, Tensor kv, Tensor! k_cache, "
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"float eps, int cache_block_size) -> ()");
|
||||
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
|
||||
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
|
||||
#endif
|
||||
|
||||
// Apply repetition penalties to logits in-place
|
||||
ops.def(
|
||||
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
|
||||
@@ -259,8 +240,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"rotary_embedding(Tensor positions, Tensor! query,"
|
||||
" Tensor!? key, int head_size,"
|
||||
" Tensor cos_sin_cache, bool is_neox, int "
|
||||
"rope_dim_offset=0, bool inverse=False) -> ()");
|
||||
" Tensor cos_sin_cache, bool is_neox) -> ()");
|
||||
ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
|
||||
|
||||
// Quantization ops
|
||||
|
||||
+4
-4
@@ -22,7 +22,7 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=13.0.2
|
||||
ARG CUDA_VERSION=13.0.0
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Explicitly set the list to avoid issues with torch 2.2
|
||||
# See https://github.com/pytorch/pytorch/pull/123243
|
||||
# From versions.json: .torch.cuda_arch_list
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
#################### BUILD BASE IMAGE ####################
|
||||
|
||||
@@ -538,7 +538,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
cuda-nvrtc-${CUDA_VERSION_DASH} \
|
||||
cuda-cuobjdump-${CUDA_VERSION_DASH} \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
@@ -765,7 +765,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# install kv_connectors if requested
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
|
||||
|
||||
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system numba==0.65.0
|
||||
uv pip install --system numba==0.61.2
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
|
||||
+5
-38
@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="bf4a7214"
|
||||
ARG RIXL_BRANCH="f33a5599"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="7009d7a1"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ARG UCX_BRANCH="da3fac2a"
|
||||
ARG UCX_REPO="https://github.com/ROCm/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
|
||||
--disable-doxygen-doc \
|
||||
--enable-optimizations \
|
||||
--enable-devel-headers \
|
||||
--with-rocm=${ROCM_PATH} \
|
||||
--with-rocm=/opt/rocm \
|
||||
--with-verbs \
|
||||
--with-dm \
|
||||
--enable-mt && \
|
||||
@@ -186,12 +186,7 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
|
||||
ninja install
|
||||
|
||||
# Generate RIXL wheel
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
RUN cd /opt/rixl && mkdir -p /app/install && \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
@@ -312,30 +307,6 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
|
||||
# Fail if git-based package dependencies are found in requirements files
|
||||
# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
|
||||
# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
|
||||
RUN echo "Checking for git-based packages in requirements files..." \
|
||||
&& echo "Checking common.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
|
||||
echo "ERROR: Git-based packages found in common.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in common.txt"; \
|
||||
fi \
|
||||
&& echo "Checking rocm.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
|
||||
echo "ERROR: Git-based packages found in rocm.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in rocm.txt"; \
|
||||
fi \
|
||||
&& echo "All requirements files are clean - no git-based packages found"
|
||||
|
||||
# Pin vLLM dependencies to exact versions of custom ROCm wheels
|
||||
# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
|
||||
COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
|
||||
@@ -436,10 +407,6 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
|
||||
ENV MIOPEN_DEBUG_CONV_DIRECT=0
|
||||
ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
|
||||
# Source code is used in the `python_only_compile.sh` test
|
||||
# We hide it inside `src/` so that this source code
|
||||
# will not be imported by other tests
|
||||
|
||||
@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.0
|
||||
uv pip install triton-xpu==3.6.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
|
||||
}
|
||||
|
||||
variable "TORCH_CUDA_ARCH_LIST" {
|
||||
default = "8.0 8.9 9.0 10.0 11.0 12.0"
|
||||
default = "8.0 8.9 9.0 10.0"
|
||||
}
|
||||
|
||||
variable "COMMIT" {
|
||||
@@ -88,6 +88,7 @@ target "test-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
@@ -99,6 +100,7 @@ target "openai-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "13.0.2"
|
||||
"default": "13.0.0"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
@@ -11,10 +11,10 @@
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
|
||||
},
|
||||
"GET_PIP_URL": {
|
||||
"default": "https://bootstrap.pypa.io/get-pip.py"
|
||||
@@ -32,7 +32,7 @@
|
||||
"default": "false"
|
||||
},
|
||||
"TORCH_CUDA_ARCH_LIST": {
|
||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
"default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
|
||||
},
|
||||
"MAX_JOBS": {
|
||||
"default": "2"
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 315 KiB After Width: | Height: | Size: 315 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 108 KiB |
File diff suppressed because one or more lines are too long
@@ -108,38 +108,6 @@ P99 ITL (ms): 8.39
|
||||
==================================================
|
||||
```
|
||||
|
||||
#### Results Visualization
|
||||
|
||||
The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--endpoint /v1/completions \
|
||||
--dataset-name sharegpt \
|
||||
--dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
|
||||
--num-prompts 100 \
|
||||
--plot-timeline \
|
||||
--timeline-itl-thresholds 2,5 \
|
||||
--plot-dataset-stats \
|
||||
--save-result
|
||||
```
|
||||
|
||||
##### Interactive Timeline
|
||||
|
||||
The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
|
||||
|
||||
Example output:
|
||||
|
||||
<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
|
||||
|
||||
##### Dataset statistics
|
||||
|
||||
The generated figure shows the input prompt and output tokens distribution.
|
||||
|
||||
Example output: 
|
||||
|
||||
#### Custom Dataset
|
||||
|
||||
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
|
||||
|
||||
@@ -155,9 +155,9 @@ switch to `--physcpubind=<cpu-list> --membind=<node>`.
|
||||
|
||||
These `--numa-bind*` options only apply to GPU execution processes. They do not
|
||||
configure the CPU backend's separate thread-affinity controls. Automatic
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as
|
||||
ROCM-based platforms; other GPU backends must provide explicit binding lists if
|
||||
they use these options.
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based platforms;
|
||||
other GPU backends must provide explicit binding lists if they use these
|
||||
options.
|
||||
|
||||
`--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in
|
||||
the same order as the GPU indices.
|
||||
|
||||
@@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model:
|
||||
|
||||
See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.
|
||||
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
|
||||
#### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)
|
||||
|
||||
@@ -75,20 +75,21 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
stt_params: SpeechToTextParams,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
) -> PromptType:
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
task_type = stt_params.task_type
|
||||
|
||||
# Example with a free-form instruction prompt
|
||||
task_word = "Transcribe" if task_type == "transcribe" else "Translate"
|
||||
prompt = (
|
||||
"<start_of_turn>user\n"
|
||||
@@ -111,22 +112,20 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
stt_params: SpeechToTextParams,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
) -> PromptType:
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
language = stt_params.language
|
||||
task_type = stt_params.task_type
|
||||
request_prompt = stt_params.request_prompt
|
||||
|
||||
if language is None:
|
||||
raise ValueError("Language must be specified")
|
||||
|
||||
@@ -214,13 +213,15 @@ Relevant server logic:
|
||||
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
|
||||
prompts = []
|
||||
for chunk in chunks:
|
||||
stt_params = request.build_stt_params(
|
||||
prompt = self.model_cls.get_generation_prompt(
|
||||
audio=chunk,
|
||||
stt_config=self.asr_config,
|
||||
model_config=self.model_config,
|
||||
language=language,
|
||||
task_type=self.task_type,
|
||||
request_prompt=request.prompt,
|
||||
to_language=to_language,
|
||||
)
|
||||
prompt = self.model_cls.get_generation_prompt(stt_params)
|
||||
prompts.append(prompt)
|
||||
return prompts, duration
|
||||
```
|
||||
|
||||
@@ -4,7 +4,6 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
|
||||
|
||||
- [Deployment with CPUs](#deployment-with-cpus)
|
||||
- [Deployment with GPUs](#deployment-with-gpus)
|
||||
- [Serving with gRPC](#serving-with-grpc)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated)
|
||||
- [Conclusion](#conclusion)
|
||||
@@ -388,49 +387,6 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
|
||||
If the service is correctly deployed, you should receive a response from the vLLM model.
|
||||
|
||||
## Serving with gRPC
|
||||
|
||||
vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies:
|
||||
|
||||
```bash
|
||||
pip install vllm[grpc]
|
||||
```
|
||||
|
||||
When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24).
|
||||
|
||||
To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes:
|
||||
|
||||
```yaml
|
||||
containers:
|
||||
- name: mistral-7b
|
||||
image: vllm/vllm-openai:latest
|
||||
command: ["/bin/sh", "-c"]
|
||||
args: [
|
||||
"pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code"
|
||||
]
|
||||
ports:
|
||||
- containerPort: 50051
|
||||
livenessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 10
|
||||
readinessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 5
|
||||
```
|
||||
|
||||
!!! note
|
||||
The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`.
|
||||
|
||||
You can also verify the health service manually with `grpcurl`:
|
||||
|
||||
```bash
|
||||
grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
|
||||
|
||||
@@ -172,7 +172,7 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
@@ -213,7 +213,7 @@ configuration.
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
|
||||
@@ -44,7 +44,7 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
|
||||
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
|
||||
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
|
||||
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
|
||||
| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
|
||||
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
|
||||
@@ -152,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt
|
||||
|
||||
- `FLASHINFER`: CUDA sm100+ with FlashInfer installed
|
||||
|
||||
`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
|
||||
`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
|
||||
|
||||
The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
|
||||
with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
|
||||
|
||||
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
|
||||
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
|
||||
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
|
||||
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
# Context Extension
|
||||
|
||||
!!! note
|
||||
The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
|
||||
This directory contains examples for extending the context length of models using vLLM.
|
||||
|
||||
## Offline Inference Example
|
||||
|
||||
The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
|
||||
### Usage
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/context_extension.py
|
||||
```
|
||||
|
||||
## OpenAI Online Method
|
||||
|
||||
You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
|
||||
|
||||
### Usage
|
||||
|
||||
Run the vLLM server with the following command to extend the context length using YARN:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
|
||||
--max-model-len 131072
|
||||
```
|
||||
|
||||
### Client Example
|
||||
|
||||
After starting the server, you can use the OpenAI Python client to interact with it:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:8000/v1",
|
||||
api_key="token-abc123" # Dummy API key, required by the client
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
max_tokens=128,
|
||||
temperature=0.8,
|
||||
top_p=0.95
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### Key Parameters
|
||||
|
||||
The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
|
||||
|
||||
Common parameters include:
|
||||
|
||||
- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
|
||||
- `factor`: The factor by which to extend the context length
|
||||
- `original_max_position_embeddings`: The original maximum position embeddings of the model
|
||||
|
||||
The following parameters are specific to vLLM:
|
||||
|
||||
- `max_model_len`: The new maximum sequence length after extension (original * factor).
|
||||
Used for KV cache pre‑allocation and request limit at serving time.
|
||||
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
|
||||
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
|
||||
|
||||
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
|
||||
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
|
||||
!!! note
|
||||
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
|
||||
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
|
||||
### Audio Inputs
|
||||
|
||||
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
|
||||
|
||||
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output:", chat_response.choices[0].message.content)
|
||||
```
|
||||
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! tip
|
||||
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
|
||||
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output from image url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
|
||||
@@ -780,70 +780,6 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
|
||||
Works with common video formats like MP4 when using OpenCV backends.
|
||||
|
||||
#### Pre-extracted Frame Sequences with `media_io_kwargs`
|
||||
|
||||
When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction.
|
||||
|
||||
**Supported Parameters:**
|
||||
|
||||
| Parameter | Type | Description |
|
||||
| --------- | ---- | ----------- |
|
||||
| `fps` | float | Frame rate of the original video |
|
||||
| `frames_indices` | list[int] | Indices of the actually sampled frames |
|
||||
| `total_num_frames` | int | Total frame count of the original video |
|
||||
| `duration` | float | Duration of the original video in seconds |
|
||||
| `do_sample_frames` | bool | Whether to perform frame sampling |
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
|
||||
|
||||
# Client-side frame extraction
|
||||
frames = extract_frames(video_path, num_frames=32)
|
||||
frames_b64 = ",".join([encode_image(f) for f in frames])
|
||||
video_url = f"data:video/jpeg;base64,{frames_b64}"
|
||||
|
||||
# Pass video metadata via media_io_kwargs
|
||||
response = client.chat.completions.create(
|
||||
model="your-multimodal-model",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "video_url", "video_url": {"url": video_url}},
|
||||
{"type": "text", "text": "Describe what happens in this video."}
|
||||
]
|
||||
}],
|
||||
extra_body={
|
||||
"media_io_kwargs": {
|
||||
"video": {
|
||||
"fps": 30.0,
|
||||
"frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90,
|
||||
100, 110, 120, 130, 140, 150, 160, 170,
|
||||
180, 190, 200, 210, 220, 230, 240, 250,
|
||||
260, 270, 280, 290, 300, 310],
|
||||
"total_num_frames": 900,
|
||||
"duration": 30.0,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
**Why use `media_io_kwargs`?**
|
||||
|
||||
When extracting frames client-side, the server loses important context about the original video:
|
||||
|
||||
- **Temporal information**: Which frames were sampled and their positions in the original timeline
|
||||
- **Video duration**: How long the original video was
|
||||
- **Frame rate**: The original playback speed
|
||||
|
||||
By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped.
|
||||
|
||||
#### Custom RGBA Background Color
|
||||
|
||||
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
|
||||
@@ -958,7 +894,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
|
||||
print("Chat completion output from audio url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
|
||||
|
||||
@@ -20,7 +20,6 @@ The following are the supported quantization formats for vLLM:
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
|
||||
|
||||
## Supported Hardware
|
||||
|
||||
|
||||
@@ -1,109 +0,0 @@
|
||||
# FP8 ViT Encoder Attention
|
||||
|
||||
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
|
||||
short text prompts/generation, the ViT encoder attention can become a significant
|
||||
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
|
||||
supports optional FP8 quantization for the ViT encoder attention via the
|
||||
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
|
||||
cuDNN attention call.
|
||||
|
||||
!!! note
|
||||
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
|
||||
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
|
||||
- Dynamic scaling is not compatible with ViT full CUDA graphs.
|
||||
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
|
||||
requests. Smaller images may see no speedup due to quantization overhead
|
||||
(3 quantization kernel launches + un-padding).
|
||||
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
|
||||
|
||||
## Requirements
|
||||
|
||||
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
|
||||
|
||||
## Usage
|
||||
|
||||
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
|
||||
with `--mm-encoder-attn-backend FLASHINFER`:
|
||||
|
||||
```bash
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
|
||||
buffer of observed Q/K/V amax values drives per-forward scale updates. This
|
||||
matches BF16 accuracy without any calibration but adds a small per-forward
|
||||
overhead.
|
||||
|
||||
## Calibrate-Once, Reuse Workflow (Recommended)
|
||||
|
||||
For production, calibrate static scales on a representative dataset once and
|
||||
reuse them to avoid the dynamic overhead:
|
||||
|
||||
```bash
|
||||
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
|
||||
# then dumps the learned scales to JSON).
|
||||
vllm bench mm-processor \
|
||||
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
|
||||
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--num-prompts 100
|
||||
|
||||
# Step 2: serve with static scales (no dynamic overhead).
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-path /path/to/scales.json
|
||||
```
|
||||
|
||||
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
|
||||
`1.5`) to leave headroom against activation outliers not present in the
|
||||
calibration set. The default has been validated to generalize across datasets
|
||||
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
|
||||
|
||||
## Scale File Format
|
||||
|
||||
```json
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
|
||||
}
|
||||
```
|
||||
|
||||
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
|
||||
|
||||
## Performance
|
||||
|
||||
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
|
||||
|
||||
| Hardware | BF16 | FP8 | Speedup |
|
||||
| -------- | ---- | ---- | ------- |
|
||||
| GB200 | 350 us | 312 us | **1.12x** |
|
||||
| GB300 | 300 us | 211 us | **1.42x** |
|
||||
|
||||
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
|
||||
|
||||
| Resolution | BF16 median | FP8 median | Speedup |
|
||||
| ---------- | ----------- | ---------- | ------- |
|
||||
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
|
||||
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
|
||||
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
|
||||
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
|
||||
|
||||
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
|
||||
|
||||
## Accuracy
|
||||
|
||||
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
|
||||
on VisionArena-Chat (with default 1.5x margin):
|
||||
|
||||
| Metric | BF16 | FP8 dynamic | FP8 static |
|
||||
| ------ | ---- | ----------- | ---------- |
|
||||
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
|
||||
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
|
||||
| exact_match | 0.584 | 0.582 | 0.578 |
|
||||
|
||||
All three configurations match within statistical noise, confirming that
|
||||
static scales calibrated on one dataset generalize to another.
|
||||
@@ -202,7 +202,7 @@ The reasoning content is also available when both tool calling and the reasoning
|
||||
print(f"Arguments: {tool_call.arguments}")
|
||||
```
|
||||
|
||||
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
|
||||
## Server-Level Default Chat Template Kwargs
|
||||
|
||||
|
||||
@@ -35,99 +35,6 @@ For reproducible measurements in your environment, use
|
||||
[`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
|
||||
or the [benchmark CLI guide](../../benchmarking/cli.md).
|
||||
|
||||
## `--speculative-config` schema
|
||||
|
||||
Use `--speculative-config` to pass speculative decoding settings as a JSON
|
||||
object on the CLI:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "draft_model",
|
||||
"model": "<draft-model>",
|
||||
"num_speculative_tokens": 5
|
||||
}'
|
||||
```
|
||||
|
||||
The same keys are accepted from Python via `LLM(..., speculative_config={...})`.
|
||||
The tables below highlight common user-facing keys accepted in this JSON
|
||||
object; they are not an exhaustive schema reference.
|
||||
For more details, see the generated [engine arguments reference](../../configuration/engine_args.md)
|
||||
and the API docs for [vllm.config.SpeculativeConfig][].
|
||||
|
||||
### Common keys
|
||||
|
||||
These keys are commonly used across speculative decoding setups, though some
|
||||
only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
|
||||
`dflash`.
|
||||
|
||||
| Key | Type | Default | Allowed values / meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
|
||||
| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
|
||||
| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
|
||||
| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
|
||||
| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. |
|
||||
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
|
||||
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
|
||||
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
|
||||
|
||||
### Method-specific keys
|
||||
|
||||
#### N-gram
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. |
|
||||
| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 4,
|
||||
"prompt_lookup_min": 2,
|
||||
"prompt_lookup_max": 5
|
||||
}'
|
||||
```
|
||||
|
||||
#### Suffix decoding
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. |
|
||||
| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. |
|
||||
| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. |
|
||||
| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 8,
|
||||
"suffix_decoding_max_tree_depth": 24,
|
||||
"suffix_decoding_max_cached_requests": 10000,
|
||||
"suffix_decoding_max_spec_factor": 1.0,
|
||||
"suffix_decoding_min_token_prob": 0.1
|
||||
}'
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--speculative-config` expects a JSON object on the CLI. In YAML config
|
||||
files, use a nested mapping instead of an escaped JSON string.
|
||||
- `tensor_parallel_size` is not a valid key in `speculative_config`. Use
|
||||
`draft_tensor_parallel_size` instead.
|
||||
- Keys such as `temperature` and `top_p` are sampling parameters, not
|
||||
`--speculative-config` fields.
|
||||
- Internal fields such as `target_model_config`, `draft_model_config`,
|
||||
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
|
||||
are populated by vLLM and are not intended to be set by users.
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
|
||||
@@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
```
|
||||
|
||||
The code used to request as completions as a client remains unchanged:
|
||||
@@ -77,8 +77,4 @@ The code used to request as completions as a client remains unchanged:
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative-config` to set all configurations related
|
||||
to speculative decoding. The previous method of specifying the model
|
||||
through `--speculative-model` and adding related parameters such as
|
||||
`--num-speculative-tokens` separately has been deprecated. For supported
|
||||
keys and examples, see the [`--speculative-config` schema](README.md#--speculative-config-schema).
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
|
||||
|
||||
@@ -38,7 +38,7 @@ for output in outputs:
|
||||
```bash
|
||||
vllm serve XiaomiMiMo/MiMo-7B-Base \
|
||||
--tensor-parallel-size 1 \
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
--speculative_config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
@@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
```
|
||||
|
||||
## Pre-trained PARD weights
|
||||
|
||||
@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--build-arg CUDA_VERSION=13.0.1 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
|
||||
--build-arg max_jobs=256 \
|
||||
--build-arg nvcc_threads=2 \
|
||||
--build-arg RUN_WHEEL_CHECK=false \
|
||||
|
||||
@@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt
|
||||
|
||||
!!! note
|
||||
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
|
||||
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
- For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
|
||||
- Finally, build and install vLLM XPU backend:
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
|
||||
| `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier |
|
||||
| `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | nan | nan |
|
||||
| `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) |
|
||||
|
||||
!!! note
|
||||
@@ -86,15 +86,14 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|
||||
### Pooling Usages
|
||||
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
|
||||
We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.
|
||||
|
||||
@@ -102,9 +101,9 @@ For more detailed information, please refer to the link below.
|
||||
|
||||
- [Classification Usages](classify.md)
|
||||
- [Embedding Usages](embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Token Classification Usages](token_classify.md)
|
||||
- [Token Embedding Usages](token_embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Scoring Usages](scoring.md)
|
||||
- [Specific Model Examples](specific_models.md)
|
||||
|
||||
@@ -114,17 +113,15 @@ Each pooling model in vLLM supports one or more of these tasks according to
|
||||
[Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
|
||||
enabling the corresponding APIs.
|
||||
|
||||
### Offline APIs corresponding to pooling usages
|
||||
### Offline APIs corresponding to pooling tasks
|
||||
|
||||
| Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function |
|
||||
|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
|
||||
| Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier |
|
||||
| Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity |
|
||||
| Token Classification Usages | N/A | `token_classify` | N/A | N/A |
|
||||
| Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) |
|
||||
| Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A |
|
||||
| Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A |
|
||||
| Plugins Usages | N/A | `plugin` | N/A | N/A |
|
||||
| Task | APIs |
|
||||
|------------------|---------------------------------------------------------------------------------------|
|
||||
| `embed` | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
|
||||
| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)` |
|
||||
| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` |
|
||||
| `token_embed` | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)` |
|
||||
| `plugin` | `LLM.encode(..., pooling_task="plugin")` |
|
||||
|
||||
!!! note
|
||||
Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
@@ -150,7 +147,7 @@ It is primarily designed for [score models](scoring.md).
|
||||
|
||||
The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
|
||||
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
|
||||
### Examples
|
||||
|
||||
@@ -186,12 +183,9 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all
|
||||
|
||||
The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
|
||||
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
|
||||
Code examples:
|
||||
|
||||
- [Online example](../../../examples/pooling/reward/token_reward_online.py)
|
||||
- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)
|
||||
Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
|
||||
|
||||
### Examples
|
||||
|
||||
|
||||
@@ -134,13 +134,3 @@ print(f"Data: {data!r}")
|
||||
## Online Serving
|
||||
|
||||
Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
|
||||
|
||||
## More examples
|
||||
|
||||
More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
|
||||
|
||||
## Deprecated Features
|
||||
|
||||
### `LLM.reward`
|
||||
|
||||
`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.
|
||||
|
||||
@@ -384,7 +384,6 @@ th {
|
||||
| `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
|
||||
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
|
||||
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
|
||||
| `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -420,7 +419,6 @@ th {
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
|
||||
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -439,7 +437,6 @@ th {
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
|
||||
@@ -475,7 +472,6 @@ th {
|
||||
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
|
||||
| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
|
||||
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
|
||||
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
|
||||
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -591,7 +587,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
@@ -646,10 +641,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
!!! note
|
||||
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
|
||||
MobileNet-v5 vision backbone.
|
||||
|
||||
|
||||
Performance is not yet fully optimized mainly due to:
|
||||
|
||||
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
|
||||
|
||||
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
|
||||
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -251,7 +251,7 @@ The following extra parameters are supported:
|
||||
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
|
||||
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
|
||||
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
|
||||
!!! note
|
||||
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
|
||||
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
|
||||
|
||||
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
|
||||
|
||||
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
|
||||
!!! note
|
||||
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
|
||||
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
|
||||
|
||||
#### Extra Parameters
|
||||
|
||||
|
||||
+2
-11
@@ -1,15 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Example offline usage of token reward models.
|
||||
|
||||
The key distinction between sequence classification and token classification
|
||||
lies in their output granularity: sequence classification produces a single
|
||||
result for an entire input sequence, whereas token classification yields a
|
||||
result for each individual token within the sequence.
|
||||
"""
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
@@ -45,14 +36,14 @@ def main(args: Namespace):
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate rewards. The output is a list of PoolingRequestOutput.
|
||||
outputs = llm.encode(prompts, pooling_task="token_classify")
|
||||
outputs = llm.reward(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
rewards = output.outputs.data
|
||||
print(f"Prompt: {prompt!r}")
|
||||
print_embeddings(rewards.tolist(), prefix="Reward")
|
||||
print_embeddings(rewards, prefix="Reward")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
Regular → Executable
+6
-6
@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
|
||||
|
||||
```bash
|
||||
# Audio + image + video
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
-q mixed_modalities
|
||||
|
||||
# Read vision and audio inputs from a single video file
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
-q use_audio_in_video
|
||||
|
||||
# Multiple audios
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
-q multi_audios
|
||||
```
|
||||
|
||||
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
|
||||
|
||||
```bash
|
||||
# Process audio inputs
|
||||
python examples/generate/multimodal/audio_language_offline.py \
|
||||
python examples/offline_inference/audio_language.py \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process image inputs
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
python examples/offline_inference/vision_language.py \
|
||||
--modality image \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process video inputs
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
python examples/offline_inference/vision_language.py \
|
||||
--modality video \
|
||||
--model-type qwen2_5_omni
|
||||
```
|
||||
Regular → Executable
+6
-36
@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
|
||||
# MiniCPM-V
|
||||
def run_minicpmv_base(questions: list[str], modality: str, model_name):
|
||||
assert modality in ["image", "video", "image+video"]
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
|
||||
|
||||
# 2.0
|
||||
# The official repo doesn't work yet, so we need to use a fork for now
|
||||
@@ -2463,12 +2463,6 @@ MODELS_NEED_VIDEO_METADATA = [
|
||||
]
|
||||
|
||||
|
||||
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
]
|
||||
|
||||
|
||||
def get_multi_modal_input(args):
|
||||
"""
|
||||
return {
|
||||
@@ -2581,29 +2575,6 @@ def apply_image_repeat(
|
||||
return inputs, inputs_with_empty_media
|
||||
|
||||
|
||||
def maybe_add_vit_cuda_graph_compilation_config(args, engine_args):
|
||||
model = args.model_type
|
||||
modality = args.modality
|
||||
enable_vit_cuda_graph = args.enable_vit_cuda_graph
|
||||
|
||||
if enable_vit_cuda_graph and model in MODELS_SUPPORT_VIT_CUDA_GRAPH:
|
||||
if modality == "image" or modality == "video":
|
||||
vision_items_per_batch = 1
|
||||
elif modality == "image+video":
|
||||
vision_items_per_batch = 2
|
||||
else:
|
||||
raise ValueError(
|
||||
f"modality={modality} is not supported for vit cuda graph."
|
||||
)
|
||||
|
||||
engine_args.compilation_config = {
|
||||
"cudagraph_mm_encoder": True,
|
||||
"encoder_cudagraph_max_vision_items_per_batch": vision_items_per_batch,
|
||||
}
|
||||
|
||||
return engine_args
|
||||
|
||||
|
||||
@contextmanager
|
||||
def time_counter(enable: bool):
|
||||
if enable:
|
||||
@@ -2654,28 +2625,33 @@ def parse_args():
|
||||
default=0,
|
||||
help="Set the seed when initializing `vllm.LLM`.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--image-repeat-prob",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Simulates the hit-ratio for multi-modal preprocessor cache (if enabled)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--disable-mm-processor-cache",
|
||||
action="store_true",
|
||||
help="If True, disables caching of multi-modal processor.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--time-generate",
|
||||
action="store_true",
|
||||
help="If True, then print the total generate() call time",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--use-different-prompt-per-request",
|
||||
action="store_true",
|
||||
help="If True, then use different prompt (with the same multi-modal "
|
||||
"data) for each request.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--verify-mm-cache-hit-with-uuids",
|
||||
action="store_true",
|
||||
@@ -2689,11 +2665,6 @@ def parse_args():
|
||||
default=None,
|
||||
help="Tensor parallel size to override the model's default setting. ",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-vit-cuda-graph",
|
||||
action="store_true",
|
||||
help="If True, will enable vit cuda graph capture and replay for the model.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -2727,7 +2698,6 @@ def main(args):
|
||||
engine_args.mm_processor_cache_gb = mm_processor_cache_gb
|
||||
if args.tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = args.tensor_parallel_size
|
||||
engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
# Don't want to check the flag multiple times, so just hijack `prompts`.
|
||||
Regular → Executable
@@ -10,8 +10,9 @@ import uuid
|
||||
|
||||
import aiohttp
|
||||
import msgpack
|
||||
import regex as re
|
||||
import zmq
|
||||
from quart import Quart, Request, make_response, request
|
||||
from quart import Quart, make_response, request
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
|
||||
MoRIIOConstants,
|
||||
@@ -24,10 +25,32 @@ decode_instances: list[dict] = []
|
||||
request_nums = 0
|
||||
app = Quart(__name__)
|
||||
|
||||
IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
|
||||
|
||||
|
||||
TRANSFER_TYPE = None
|
||||
|
||||
|
||||
def _append_whole_dict_unique(target_list, data_dict):
|
||||
new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
|
||||
for existed in target_list:
|
||||
existed_filtered = {k: v for k, v in existed.items() if k != "index"}
|
||||
if existed_filtered == new_filtered:
|
||||
return False
|
||||
print("!!APPEND!!", data_dict)
|
||||
target_list.append(data_dict)
|
||||
transfer_mode = data_dict.get("transfer_mode", "unknown")
|
||||
global TRANSFER_TYPE
|
||||
|
||||
if TRANSFER_TYPE is None:
|
||||
TRANSFER_TYPE = transfer_mode
|
||||
logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
|
||||
elif transfer_mode != TRANSFER_TYPE:
|
||||
raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
_list_lock = threading.RLock()
|
||||
|
||||
|
||||
@@ -45,81 +68,23 @@ def _listen_for_register(hostname, port):
|
||||
if router_socket in socks:
|
||||
remote_addr, msg = router_socket.recv_multipart()
|
||||
data = msgpack.loads(msg)
|
||||
if data.get("type") == "HELLO":
|
||||
if data["type"] == "HELLO":
|
||||
pass
|
||||
elif data.get("type") in ("P", "D"):
|
||||
role = data["type"]
|
||||
required_keys = {
|
||||
"http_address",
|
||||
"zmq_address",
|
||||
"dp_size",
|
||||
"tp_size",
|
||||
"transfer_mode",
|
||||
}
|
||||
missing = required_keys - data.keys()
|
||||
if missing:
|
||||
logger.error(
|
||||
"Registration message missing required keys %s; skipping",
|
||||
missing,
|
||||
)
|
||||
continue
|
||||
# Derive request_address from http_address
|
||||
# api path suffix is appended at request time
|
||||
instance = {
|
||||
"role": role,
|
||||
"request_address": f"http://{data['http_address']}/v1",
|
||||
"http_address": data["http_address"],
|
||||
"zmq_address": data["zmq_address"],
|
||||
"dp_size": data["dp_size"],
|
||||
"tp_size": data["tp_size"],
|
||||
"transfer_mode": data["transfer_mode"],
|
||||
}
|
||||
# zmq_address format: "host:IP,handshake:PORT,notify:PORT"
|
||||
# Stored verbatim; embedded into the request_id by handle_request.
|
||||
|
||||
global TRANSFER_TYPE
|
||||
transfer_mode = instance["transfer_mode"]
|
||||
target_list = prefill_instances if role == "P" else decode_instances
|
||||
elif (
|
||||
data["type"] == "register"
|
||||
and data["role"] == "P"
|
||||
and data["request_address"] not in prefill_instances
|
||||
):
|
||||
with _list_lock:
|
||||
if TRANSFER_TYPE is None:
|
||||
TRANSFER_TYPE = transfer_mode
|
||||
logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
|
||||
elif transfer_mode != TRANSFER_TYPE:
|
||||
logger.error(
|
||||
"Mismatched transfer mode: expected %s, got %s;"
|
||||
" skipping registration of %s",
|
||||
TRANSFER_TYPE,
|
||||
transfer_mode,
|
||||
data["http_address"],
|
||||
)
|
||||
continue
|
||||
existing_idx = next(
|
||||
(
|
||||
idx
|
||||
for idx, i in enumerate(target_list)
|
||||
if i.get("http_address") == data["http_address"]
|
||||
),
|
||||
None,
|
||||
)
|
||||
if existing_idx is not None:
|
||||
target_list[existing_idx] = instance
|
||||
logger.info(
|
||||
"Updated existing %s instance: %s",
|
||||
"Prefill" if role == "P" else "Decode",
|
||||
instance,
|
||||
)
|
||||
else:
|
||||
target_list.append(instance)
|
||||
logger.info(
|
||||
"Registered %s instance: %s",
|
||||
"Prefill" if role == "P" else "Decode",
|
||||
instance,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Received message with unrecognized type %r; ignoring",
|
||||
data.get("type"),
|
||||
)
|
||||
_append_whole_dict_unique(prefill_instances, data)
|
||||
|
||||
elif (
|
||||
data["type"] == "register"
|
||||
and data["role"] == "D"
|
||||
and data["request_address"] not in decode_instances
|
||||
):
|
||||
with _list_lock:
|
||||
_append_whole_dict_unique(decode_instances, data)
|
||||
|
||||
|
||||
def start_service_discovery(hostname, port):
|
||||
@@ -136,7 +101,7 @@ def start_service_discovery(hostname, port):
|
||||
|
||||
|
||||
async def send_request_to_prefill(
|
||||
endpoint, req_data, request_id, selected_prefill_dp_rank
|
||||
endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
|
||||
):
|
||||
req_data_copy = req_data
|
||||
|
||||
@@ -144,8 +109,12 @@ async def send_request_to_prefill(
|
||||
{
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_handshake_port": d_endpoint["handshake_port"],
|
||||
"remote_notify_port": d_endpoint["notify_port"],
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": dip,
|
||||
"remote_port": dport,
|
||||
}
|
||||
)
|
||||
req_data_copy["stream"] = False
|
||||
@@ -170,13 +139,10 @@ async def send_request_to_prefill(
|
||||
return await response.json()
|
||||
|
||||
else:
|
||||
error_message = (
|
||||
f"send_request_to_prefill response ={response},"
|
||||
f"reason={response.reason}, status={response.status},"
|
||||
f"method={response.method}, url={response.url},"
|
||||
f"real_url={response.real_url}"
|
||||
raise RuntimeError(
|
||||
"send_request_to_prefill response.status != 200response.status = ",
|
||||
response.status,
|
||||
)
|
||||
raise RuntimeError(error_message)
|
||||
|
||||
|
||||
async def start_decode_request(endpoint, req_data, request_id):
|
||||
@@ -197,13 +163,9 @@ async def stream_decode_response(session, response, request_id):
|
||||
async for chunk_bytes in response.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
else:
|
||||
error_message = (
|
||||
f"stream_decode_response response ={response},"
|
||||
f"reason={response.reason}, status={response.status},"
|
||||
f"method={response.method}, url={response.url},"
|
||||
f"real_url={response.real_url}"
|
||||
raise RuntimeError(
|
||||
f"decode response.status != 200, status = {response.status}"
|
||||
)
|
||||
raise RuntimeError(error_message)
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
@@ -213,22 +175,21 @@ def example_round_robin_dp_loader(request_number, dp_size):
|
||||
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
async def handle_completions_request():
|
||||
return await handle_request("/completions", request)
|
||||
|
||||
|
||||
@app.route("/v1/chat/completions", methods=["POST"])
|
||||
async def handle_chat_completions_request():
|
||||
return await handle_request("/chat/completions", request)
|
||||
|
||||
|
||||
async def handle_request(api: str, request: Request):
|
||||
async def handle_request():
|
||||
try:
|
||||
with _list_lock:
|
||||
global request_nums
|
||||
request_nums += 1
|
||||
|
||||
def extract_ip_port_fast(url):
|
||||
match = IP_PORT_PATTERN.search(url)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid URL format: {url}")
|
||||
return match.groups()
|
||||
|
||||
req_data = await request.get_json()
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
prefill_instance_endpoint = None
|
||||
decode_instance_endpoint = None
|
||||
@@ -254,14 +215,7 @@ async def handle_request(api: str, request: Request):
|
||||
prefill_instance_endpoint["dp_size"],
|
||||
)
|
||||
|
||||
# Embed both zmq_addresses in the request_id so the connector can parse
|
||||
# the peer's host/ports from it, similar to P2P-NCCL
|
||||
uid = str(uuid.uuid4()).replace("-", "")
|
||||
request_id = (
|
||||
f"___prefill_addr_{prefill_instance_endpoint['zmq_address']}"
|
||||
f"___decode_addr_{decode_instance_endpoint['zmq_address']}"
|
||||
f"_{uid}"
|
||||
)
|
||||
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
|
||||
|
||||
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
|
||||
|
||||
@@ -276,36 +230,40 @@ async def handle_request(api: str, request: Request):
|
||||
)
|
||||
req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
prefill_request_url = prefill_instance_endpoint["request_address"] + api
|
||||
send_prefill_task = asyncio.create_task(
|
||||
send_request_to_prefill(
|
||||
prefill_request_url,
|
||||
prefill_instance_endpoint["request_address"],
|
||||
req_data_to_prefill,
|
||||
request_id,
|
||||
decode_instance_endpoint,
|
||||
dip,
|
||||
dport,
|
||||
selected_prefill_dp_rank,
|
||||
)
|
||||
)
|
||||
ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
|
||||
|
||||
req_data["max_tokens"] -= 1
|
||||
|
||||
req_data["kv_transfer_params"] = {
|
||||
"do_remote_decode": False,
|
||||
"do_remote_prefill": True,
|
||||
"remote_handshake_port": prefill_instance_endpoint["handshake_port"],
|
||||
"remote_notify_port": prefill_instance_endpoint["notify_port"],
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"transfer_id": transfer_id,
|
||||
"remote_host": ip,
|
||||
"remote_port": port,
|
||||
}
|
||||
if TRANSFER_TYPE == "READ":
|
||||
# In read mode, prefill and decode are executed serially.
|
||||
prefill_response = await send_prefill_task
|
||||
prefill_kv = prefill_response["kv_transfer_params"]
|
||||
req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[
|
||||
"remote_engine_id"
|
||||
]
|
||||
req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[
|
||||
"remote_block_ids"
|
||||
]
|
||||
req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
|
||||
req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
|
||||
"kv_transfer_params"
|
||||
]["remote_engine_id"]
|
||||
req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
|
||||
"kv_transfer_params"
|
||||
]["remote_block_ids"]
|
||||
|
||||
req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
|
||||
"dp_size"
|
||||
@@ -316,10 +274,12 @@ async def handle_request(api: str, request: Request):
|
||||
|
||||
if selected_prefill_dp_rank is not None:
|
||||
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
|
||||
req_data["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
decode_request_url = decode_instance_endpoint["request_address"] + api
|
||||
decode_request_task = asyncio.create_task(
|
||||
start_decode_request(decode_request_url, req_data, request_id)
|
||||
start_decode_request(
|
||||
decode_instance_endpoint["request_address"], req_data, request_id
|
||||
)
|
||||
)
|
||||
|
||||
session, decode_response = await decode_request_task
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user