Compare commits

..
Author SHA1 Message Date
Woosuk Kwon 5c7ecf1586 Support topk 1024
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 17:34:44 +00:00
Woosuk Kwon 83c930b915 fix plan caching
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 17:31:27 +00:00
Woosuk Kwon 02473af4df cleanup
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 09:06:08 +00:00
Woosuk Kwon 667acb4917 fix test
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 09:06:08 +00:00
Woosuk Kwon 508b4719f1 Fix & simplify integration
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 08:54:28 +00:00
Woosuk Kwon d97a04434d Update
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 06:54:00 +00:00
Woosuk Kwon 202e2c7cd0 Port Fast Top-K kernel
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-26 06:53:35 +00:00
Yongye ZhuandGitHub e5108f7443 CI Failure for deep_gemm and layernorm_fp8_quant
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-04-25 23:46:58 -07:00
Yifan QiaoandGitHub 6a2e1edf98 Merge branch 'main' into feat/dsv4-support 2026-04-25 22:51:50 -07:00
36992a0fdd [Bugfix][CI] Run mooncake HMA worker tests on GPU lane (#241)
Co-authored-by: Zhewen Li <zhewenli@inferact.ai>
2026-04-26 05:20:11 +00:00
Yifan Qiao d95a973b21 fix (ci): misc api mismatches
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-26 03:41:52 +00:00
Jee Jee LiandYifan Qiao 6fac86c362 Add model information
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
2026-04-26 02:14:51 +00:00
Yifan Qiao 618e3b60da fix (ci): interface mismatches
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-25 20:04:05 +00:00
Yifan Qiao b35352718c chore: fix pre-commit
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-25 20:04:05 +00:00
qizixiandYifan Qiao f704cf3218 [Bugfix] Flatten DeepSeek V32 indexer next_n on non-SM100 archs
Signed-off-by: qizixi <zixi@inferact.ai>
2026-04-25 20:01:56 +00:00
Woosuk KwonandYifan Qiao 9abe2bdd18 free up unused weights and support dummy weights
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-25 20:01:43 +00:00
Woosuk KwonandYifan Qiao 5e3525c0c9 Integrate MegaMoE kernel
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-04-25 20:01:29 +00:00
Yifan Qiao c75c382844 fix: config
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-25 20:01:28 +00:00
Yifan Qiao cf3e4173d1 fix: update cuda requirements
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-25 20:01:28 +00:00
Yifan Qiao 908ab01672 chore: pass mypy
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-04-25 20:01:28 +00:00
+6 434d934194 feat: support deepseek v4
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yongye Zhu <yongye@inferact.ai>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Simon Mo <simon@inferact.ai>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Roy Wang <yasong.wang@inferact.ai>
Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: youkaichao <youkaichao@gmail.com>
Co-authored-by: Zhewen Li <jerven.vllm@gmail.com>
Co-authored-by: Zijing Liu <liuzijing2014@gmail.com>
Co-authored-by: khluu <khluu000@gmail.com>
Co-authored-by: qizixi <zixi@inferact.ai>
2026-04-25 20:01:28 +00:00
468 changed files with 8665 additions and 23435 deletions
+3 -14
View File
@@ -69,11 +69,11 @@ steps:
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
pytest -x -v -s tests/quantization/test_cpu_wna16.py"
- label: CPU-Distributed Tests (PP+TP)
- label: CPU-Distributed Tests
depends_on: []
device: intel_cpu
no_plugin: true
source_file_dependencies: &cpu_distributed_deps
source_file_dependencies:
- csrc/cpu/shm.cpp
- vllm/v1/worker/cpu_worker.py
- vllm/v1/worker/gpu_worker.py
@@ -82,21 +82,10 @@ steps:
- vllm/platforms/cpu.py
- vllm/distributed/parallel_state.py
- vllm/distributed/device_communicators/cpu_communicator.py
- .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh tp_pp"
- label: CPU-Distributed Tests (DP+TP)
depends_on: []
device: intel_cpu
no_plugin: true
source_file_dependencies: *cpu_distributed_deps
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh dp_tp"
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
- label: CPU-Multi-Modal Model Tests %N
depends_on: []
-1
View File
@@ -192,7 +192,6 @@ export BUILDKITE_COMMIT
export PARENT_COMMIT
export IMAGE_TAG
export IMAGE_TAG_LATEST
export COMMIT="${COMMIT:-${BUILDKITE_COMMIT}}"
export CACHE_FROM
export CACHE_FROM_BASE_BRANCH
export CACHE_FROM_MAIN
+1
View File
@@ -126,4 +126,5 @@ steps:
'cd tests &&
pytest -v -s lora/test_default_mm_loras.py &&
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
(pytest -v -s lora/test_qwenvl.py || true) &&
pytest -v -s lora/test_whisper.py'
+14 -114
View File
@@ -27,7 +27,7 @@ steps:
- "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 ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
env:
DOCKER_BUILDKIT: "1"
@@ -40,7 +40,7 @@ steps:
- "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 ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -53,7 +53,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -66,7 +66,7 @@ steps:
- "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 ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
env:
DOCKER_BUILDKIT: "1"
@@ -79,7 +79,7 @@ steps:
- "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 ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -92,7 +92,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
- "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"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -121,19 +121,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
--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 \
--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=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 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"
@@ -146,19 +134,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
--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 \
--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=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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
- label: "Build release image - x86_64 - CUDA 12.9"
@@ -168,18 +144,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
--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 \
--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 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"
# 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"
@@ -192,18 +157,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
--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 \
--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 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"
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
@@ -213,21 +167,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--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 \
--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=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 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"
@@ -239,21 +179,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--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 \
--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=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 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"
@@ -263,20 +189,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
--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 \
--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 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"
@@ -288,20 +201,7 @@ steps:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
--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 \
--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 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"
- block: "Build release image for x86_64 CPU"
-142
View File
@@ -1,142 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Detect the manylinux platform tag for a wheel and rename it in place.
vLLM's build images produce wheels with the generic ``linux_<arch>`` platform
tag, which installers like ``pip`` won't accept off PyPI/our index. We need to
rewrite the platform tag to the appropriate ``manylinux_<major>_<minor>_<arch>``
before uploading.
Historically the tag was hard-coded per build (``manylinux_2_31`` for the
Ubuntu 20.04-based image, ``manylinux_2_35`` for the Ubuntu 22.04-based
images). That is brittle: bumping the base image silently produces wheels
labelled with the wrong glibc requirement. This script asks ``auditwheel``
to derive the tag from the symbol versions actually referenced by the
binaries inside the wheel, so the label tracks reality.
We can't simply call ``auditwheel repair`` -- it tries to graft external
shared libraries into the wheel and fails on vLLM's CUDA/cuBLAS dependencies.
Instead we use ``auditwheel.wheel_abi.analyze_wheel_abi`` directly, which is
the same call that powers ``auditwheel show``, and read off
``winfo.sym_policy.name``.
Usage:
detect-manylinux-tag.py <wheel_path>
The wheel is renamed in place; the new path is printed on stdout. All
diagnostics go to stderr so callers can capture stdout safely.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from auditwheel.error import (
AuditwheelError,
NonPlatformWheelError,
WheelToolsError,
)
from auditwheel.wheel_abi import analyze_wheel_abi
from auditwheel.wheeltools import get_wheel_architecture, get_wheel_libc
def detect_platform_tag(wheel_path: Path) -> str:
"""Return the most precise platform tag the wheel is consistent with.
Mirrors ``auditwheel show`` but returns ``sym_policy`` rather than
``overall_policy``: we only care about the glibc symbol versions used,
not about other policy axes (ISA extensions, blacklist, etc.) that
``overall_policy`` folds in.
"""
fn = wheel_path.name
try:
arch = get_wheel_architecture(fn)
except (WheelToolsError, NonPlatformWheelError):
# Architecture isn't deducible from the filename; let auditwheel
# infer it from the ELF binaries inside the wheel.
arch = None
try:
libc = get_wheel_libc(fn)
except WheelToolsError:
# An unrepaired wheel uses ``linux_<arch>``, which doesn't encode
# libc. Let auditwheel infer it from the ELF binaries.
libc = None
winfo = analyze_wheel_abi(
libc,
arch,
wheel_path,
frozenset(),
disable_isa_ext_check=False,
allow_graft=False,
)
return winfo.sym_policy.name
def rename_wheel(wheel_path: Path, new_platform_tag: str) -> Path:
"""Rename the wheel in place, replacing only its platform tag."""
# Wheel filename per PEP 427:
# {distribution}-{version}(-{build})?-{python}-{abi}-{platform}.whl
# The platform tag is always the last ``-``-separated token before
# ``.whl``. Compound tags like ``manylinux_2_31_x86_64`` use ``_`` as the
# internal separator, so ``-``-splitting is unambiguous.
parts = wheel_path.stem.split("-")
if len(parts) < 5:
raise ValueError(f"Unrecognised wheel filename: {wheel_path.name}")
parts[-1] = new_platform_tag
new_path = wheel_path.with_name("-".join(parts) + ".whl")
if new_path != wheel_path:
wheel_path.rename(new_path)
return new_path
def main() -> int:
parser = argparse.ArgumentParser(
description="Detect a wheel's manylinux platform tag with "
"auditwheel and rename the wheel in place."
)
parser.add_argument(
"wheel",
type=Path,
help="Path to the wheel to inspect and rename.",
)
args = parser.parse_args()
wheel_path: Path = args.wheel
if not wheel_path.is_file():
print(f"error: {wheel_path} is not a file", file=sys.stderr)
return 1
# Catch the things that ``analyze_wheel_abi`` and ``rename_wheel`` can
# raise: any subclass of ``AuditwheelError`` (pure-Python wheels,
# invalid libc, malformed wheels), filesystem errors, or our own
# ``ValueError`` for an unrecognised wheel filename. Print a single
# ``ERROR_TYPE: message`` line to stderr instead of a Python
# traceback, which is much friendlier in CI logs.
try:
new_tag = detect_platform_tag(wheel_path)
print(f"detected platform tag: {new_tag}", file=sys.stderr)
new_path = rename_wheel(wheel_path, new_tag)
except (AuditwheelError, ValueError, OSError) as e:
print(
f"error: failed to retag {wheel_path.name}: {type(e).__name__}: {e}",
file=sys.stderr,
)
return 2
if new_path != wheel_path:
print(f"renamed {wheel_path.name} -> {new_path.name}", file=sys.stderr)
else:
print(f"wheel already tagged {new_tag}", file=sys.stderr)
print(new_path)
return 0
if __name__ == "__main__":
sys.exit(main())
@@ -1,54 +0,0 @@
#!/bin/bash
# Emit docker build flags for release image provenance metadata.
# Keep this helper best-effort: missing Buildkite metadata should fall back to
# local/default values instead of blocking the Docker build.
# Variant examples: "", "cu129", "ubuntu2404", "cu129-ubuntu2404".
variant="${1:-}"
variant_suffix="${variant:+-${variant}}"
image_name="${VLLM_DOCKER_IMAGE_NAME:-vllm/vllm-openai}"
staging_repo="${VLLM_STAGING_IMAGE_REPO:-public.ecr.aws/q9t5s3a7/vllm-release-repo}"
build_commit="${VLLM_BUILD_COMMIT:-${BUILDKITE_COMMIT:-unknown}}"
build_pipeline="${VLLM_BUILD_PIPELINE:-${BUILDKITE_PIPELINE_ID:-${BUILDKITE_PIPELINE_SLUG:-local}}}"
build_url="${VLLM_BUILD_URL:-${BUILDKITE_BUILD_URL:-}}"
tag_commit="${BUILDKITE_COMMIT:-${build_commit}}"
if [[ -n "${BUILDKITE:-}" || -n "${BUILDKITE_COMMIT:-}" ]]; then
release_version="${RELEASE_VERSION:-}"
if command -v buildkite-agent >/dev/null 2>&1; then
release_version="${release_version:-$(buildkite-agent meta-data get release-version 2>/dev/null)}"
fi
release_version="${release_version#v}"
release_version="${release_version:-${tag_commit}}"
staging_image_ref="${staging_repo}:${tag_commit}-$(uname -m)${variant_suffix}"
if [[ "${NIGHTLY:-}" == "1" ]]; then
if [[ -z "${variant}" ]]; then
image_tag="${image_name}:nightly-${tag_commit}"
elif [[ "${variant}" == cu* ]]; then
cuda_variant="${variant%%-*}"
remaining_variant="${variant#${cuda_variant}}"
image_tag="${image_name}:${cuda_variant}-nightly-${tag_commit}${remaining_variant}"
else
image_tag="${image_name}:nightly-${tag_commit}${variant_suffix}"
fi
else
image_tag="${image_name}:v${release_version}${variant_suffix}"
fi
else
image_tag="${VLLM_IMAGE_TAG:-local/vllm-openai:dev}"
staging_image_ref="${image_tag}"
fi
emit_arg() {
printf -- "--build-arg %s=%s " "$1" "$2"
}
emit_arg VLLM_BUILD_COMMIT "${build_commit}"
emit_arg VLLM_BUILD_PIPELINE "${build_pipeline}"
emit_arg VLLM_BUILD_URL "${build_url}"
# This is the intended public tag. The final digest is only known after push.
emit_arg VLLM_IMAGE_TAG "${image_tag}"
printf -- "--tag %s " "${staging_image_ref}"
@@ -10,13 +10,20 @@ set -ex
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
DEFAULT_VARIANT_ALIAS="cu130" # 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/"
# Select python3 (>= 3.12) -- local if available, else a docker fallback.
# shellcheck source=lib/select-python.sh
source .buildkite/scripts/lib/select-python.sh
select_python
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
@@ -3,37 +3,42 @@ set -euox pipefail
export VLLM_CPU_CI_ENV=0
export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM
MODE=${1:-all}
echo "--- PP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename tp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/tp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
run_scenario() {
local label="$1" result_file="$2"
shift 2
echo "--- $label"
vllm serve meta-llama/Llama-3.2-3B-Instruct "$@" --max-model-len=4096 &
local server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename "$result_file" \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM "$server_pid"; wait "$server_pid" || true
if [ "$(jq '.failed' "./test_results/$result_file")" -ne 0 ]; then
echo "Some requests were failed in $label!"
exit 1
fi
}
case "$MODE" in
tp_pp) run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 ;;
dp_tp) run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 ;;
all)
run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2
run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2
;;
*) echo "ERROR: unknown mode '$MODE' (expected: tp_pp | dp_tp | all)" >&2; exit 1 ;;
esac
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
-127
View File
@@ -1,127 +0,0 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Shared helper for rewriting a wheel's platform tag from the generic
# ``linux_<arch>`` to the correct ``manylinux_<major>_<minor>_<arch>``.
# After sourcing, call ``apply_manylinux_tag <wheel>`` on each wheel
# that still carries the generic tag; the renamed path is printed on
# stdout (logs go to stderr).
#
# Why a pinned Docker container instead of using whatever Python
# happens to be on the agent:
# - vLLM's release agents are heterogeneous -- they don't agree on
# a Python minor version, and we can't rely on a particular
# ``auditwheel`` being installed.
# - ``detect-manylinux-tag.py`` reads ``auditwheel.wheel_abi`` and
# ``Policy.sym_policy``, which are *internal* APIs without a
# stability promise. Pinning both Python and auditwheel makes the
# detected tag a function of the inputs alone, and shifts version
# bumps from "implicit drift" to "deliberate, retested change".
# - Other release scripts (``generate-and-upload-nightly-index.sh``,
# ``upload-rocm-wheels.sh``) already use the python:3-slim image
# when the agent's interpreter is too old; this is the same idea
# made stricter.
#
# To keep the per-wheel cost down (the ROCm upload retags ~10 wheels
# each run), we install auditwheel into a long-lived helper container
# once on source, then ``docker exec`` into it for each call.
#
# Trap behaviour:
# - Sourcing installs an EXIT trap that calls ``manylinux_cleanup`` to
# tear down the helper container. Any EXIT trap that was already in
# place when this file was sourced is captured and run AFTER our
# cleanup, so we don't silently clobber it.
# - If a caller sets a new EXIT trap *after* sourcing, that trap will
# replace ours; in that case the caller should call
# ``manylinux_cleanup`` from their own handler.
if [[ -n "${_MANYLINUX_LIB_SOURCED:-}" ]]; then
return 0
fi
_MANYLINUX_LIB_SOURCED=1
# Pin both sides. Bump these deliberately and re-run a representative
# wheel from each build target through the detection.
_MANYLINUX_PYTHON_IMAGE="python:3.12-slim"
_MANYLINUX_AUDITWHEEL_VERSION="6.6.0"
# Resolve our own directory (and the sibling detect script) using the
# canonical, symlink-resolved path. The container mounts cwd at the
# same absolute path on both sides, so all paths we hand to it -- the
# script, the wheel -- must canonicalise to a location under cwd.
_MANYLINUX_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
_MANYLINUX_DETECT_SCRIPT="$(cd "${_MANYLINUX_LIB_DIR}/.." && pwd -P)/detect-manylinux-tag.py"
_MANYLINUX_CWD="$(pwd -P)"
docker pull --quiet "$_MANYLINUX_PYTHON_IMAGE" >/dev/null
# Spin up a long-lived helper container so we install auditwheel once
# and then ``docker exec`` into it for each wheel.
#
# The container runs as root so ``pip install`` can write into the
# system site-packages; individual ``docker exec`` calls below pin
# themselves to the host UID so any file rename happens with host
# ownership, not root.
_MANYLINUX_CONTAINER="$(docker run -d --rm \
-v "$_MANYLINUX_CWD:$_MANYLINUX_CWD" \
-w "$_MANYLINUX_CWD" \
"$_MANYLINUX_PYTHON_IMAGE" \
sleep infinity)"
docker exec "$_MANYLINUX_CONTAINER" \
pip install --quiet --disable-pip-version-check \
--root-user-action=ignore \
"auditwheel==${_MANYLINUX_AUDITWHEEL_VERSION}"
# Public cleanup -- safe to call multiple times.
manylinux_cleanup() {
if [[ -n "${_MANYLINUX_CONTAINER:-}" ]]; then
docker rm -f "$_MANYLINUX_CONTAINER" >/dev/null 2>&1 || true
_MANYLINUX_CONTAINER=""
fi
}
# Capture any EXIT trap that was already in place so we can chain to
# it rather than overwrite it. ``trap -p EXIT`` prints the handler in
# eval-able form (``trap -- 'CMD' EXIT``) or nothing if unset; we
# strip the wrapper to recover ``CMD``. Handles the common case --
# CMDs without embedded single quotes -- and degrades gracefully (we
# still run our own cleanup) for the pathological case.
_manylinux_prev_exit_trap_cmd=""
_manylinux_existing_exit_trap="$(trap -p EXIT)"
if [[ -n "$_manylinux_existing_exit_trap" ]]; then
_tmp="${_manylinux_existing_exit_trap#trap -- \'}"
_manylinux_prev_exit_trap_cmd="${_tmp%\' EXIT}"
unset _tmp
fi
unset _manylinux_existing_exit_trap
_manylinux_run_exit_chain() {
manylinux_cleanup
if [[ -n "$_manylinux_prev_exit_trap_cmd" ]]; then
eval "$_manylinux_prev_exit_trap_cmd"
fi
}
trap _manylinux_run_exit_chain EXIT
# Detect the manylinux platform tag for a single wheel and rename it
# in place, printing the renamed wheel path on stdout. Returns
# non-zero on failure (which under ``set -e`` propagates to caller).
#
# The wheel must be reachable via a path under the host cwd so it's
# visible inside the helper container; in CI the wheels always live
# under ``artifacts/`` so this is fine.
apply_manylinux_tag() {
local wheel="$1"
local abs_wheel
abs_wheel="$(realpath "$wheel")"
local new_wheel
new_wheel="$(docker exec -u "$(id -u):$(id -g)" \
"$_MANYLINUX_CONTAINER" \
python "$_MANYLINUX_DETECT_SCRIPT" "$abs_wheel")"
if [[ -z "$new_wheel" || ! -f "$new_wheel" ]]; then
echo "apply_manylinux_tag: detect-manylinux-tag.py did not produce a valid wheel path for $wheel" >&2
return 1
fi
printf '%s\n' "$new_wheel"
}
-41
View File
@@ -1,41 +0,0 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Pick a Python interpreter for buildkite scripts: prefer a local
# ``python3`` if it is recent enough (>= 3.12), otherwise fall back to
# a one-shot Docker container running ``python:3-slim``. After
# ``select_python`` returns, ``$PYTHON`` is set in the caller's shell
# and is safe to use as a command (e.g. ``$PYTHON some_script.py``).
#
# The 3.12 threshold matches what the existing nightly-index work
# expects -- typing features used by ``generate-nightly-index.py``.
# This helper does not pin the *minor* version; if you need stricter
# reproducibility (e.g. relying on auditwheel internals), invoke
# Docker yourself with a pinned tag rather than calling this.
if [[ -n "${_SELECT_PYTHON_LIB_SOURCED:-}" ]]; then
return 0
fi
_SELECT_PYTHON_LIB_SOURCED=1
# Sets ``PYTHON`` in the caller's shell and exports it. Idempotent --
# calling twice is safe and the second call simply re-runs the probe.
select_python() {
local py="${PYTHON_PROG:-python3}"
local has_new_python
has_new_python=$("$py" -c \
"print(1 if __import__('sys').version_info >= (3,12) else 0)" \
2>/dev/null || echo 0)
if [[ "$has_new_python" -eq 0 ]]; then
# ``-u $(id -u):$(id -g)`` so files created via the container
# end up owned by the host user, not root.
docker pull python:3-slim
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
else
PYTHON="$py"
fi
export PYTHON
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
}
@@ -28,7 +28,6 @@
# BFCL_MAX_MODEL_LEN - Max model length (default: 4096)
# BFCL_PORT - Server port (default: 8000)
# BFCL_REASONING_PARSER - Reasoning parser name (default: disabled)
# BFCL_TEMPERATURE - Temperature (default: 0.0)
# BFCL_EXTRA_ARGS - Additional vLLM server args
set -euo pipefail
@@ -44,7 +43,6 @@ TP_SIZE="${BFCL_TP_SIZE:-1}"
MAX_MODEL_LEN="${BFCL_MAX_MODEL_LEN:-4096}"
PORT="${BFCL_PORT:-8000}"
REASONING_PARSER="${BFCL_REASONING_PARSER:-}"
TEMPERATURE="${BFCL_TEMPERATURE:-0.0}"
EXTRA_ARGS="${BFCL_EXTRA_ARGS:-}"
# Set up output directory
@@ -141,7 +139,7 @@ echo "vLLM server is ready. (started in ${SECONDS_WAITED}s)"
# be patched in-process so BFCL knows to use the OpenAI-compatible handler
# against our local vLLM server.
bfcl_exit_code=0
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$TEMPERATURE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
import os
import sys
@@ -150,8 +148,7 @@ test_category = sys.argv[2]
num_threads = int(sys.argv[3])
port = sys.argv[4]
api_type = sys.argv[5]
temperature = float(sys.argv[6])
output_dir = sys.argv[7] if len(sys.argv) > 7 and sys.argv[7] else os.getcwd()
output_dir = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else os.getcwd()
os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
os.environ["OPENAI_API_KEY"] = "dummy"
@@ -207,7 +204,6 @@ gen_kwargs["model"] = [model]
gen_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
gen_kwargs["skip_server_setup"] = True
gen_kwargs["num_threads"] = num_threads
gen_kwargs["temperature"] = temperature
generate(**gen_kwargs)
# ---- evaluate ----
+14 -8
View File
@@ -2,18 +2,14 @@
set -ex
# Upload a single wheel to S3, after detecting and applying the appropriate
# manylinux platform tag with auditwheel.
# Upload a single wheel to S3 (rename linux -> manylinux).
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
# shellcheck source=lib/manylinux.sh
source .buildkite/scripts/lib/manylinux.sh
BUCKET="vllm-wheels"
SUBPATH=$BUILDKITE_COMMIT
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# ========= locate the wheel ==========
# ========= collect, rename & upload the wheel ==========
# Assume wheels are in artifacts/dist/*.whl
wheel_files=(artifacts/dist/*.whl)
@@ -25,9 +21,19 @@ if [[ ${#wheel_files[@]} -ne 1 ]]; then
fi
wheel="${wheel_files[0]}"
# ========= detect manylinux tag and rename ==========
# default build image uses ubuntu 20.04, which corresponds to manylinux_2_31
# we also accept params as manylinux tag
# refer to https://github.com/mayeut/pep600_compliance?tab=readme-ov-file#acceptable-distros-to-build-wheels
manylinux_version="${1:-manylinux_2_31}"
wheel="$(apply_manylinux_tag "$wheel")"
# Rename 'linux' to the appropriate manylinux version in the wheel filename
if [[ "$wheel" != *"linux"* ]]; then
echo "Error: Wheel filename does not contain 'linux': $wheel"
exit 1
fi
new_wheel="${wheel/linux/$manylinux_version}"
mv -- "$wheel" "$new_wheel"
wheel="$new_wheel"
echo "Renamed wheel to: $wheel"
# Extract the version from the wheel
+18 -23
View File
@@ -20,6 +20,10 @@ BUCKET="${S3_BUCKET:-vllm-wheels}"
ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
INDICES_OUTPUT_DIR="rocm-indices"
PYTHON="${PYTHON_PROG:-python3}"
# ROCm uses manylinux_2_35 (Ubuntu 22.04 based)
MANYLINUX_VERSION="manylinux_2_35"
echo "========================================"
echo "ROCm Wheel Upload Configuration"
@@ -30,21 +34,19 @@ echo "Commit: $BUILDKITE_COMMIT"
echo "Branch: $BUILDKITE_BRANCH"
echo "========================================"
# ======== Part 0: Setup Python and helpers ========
# ======== Part 0: Setup Python ========
# Pick a Python interpreter for index generation -- local if recent
# enough, else a one-shot docker fallback.
# shellcheck source=lib/select-python.sh
source .buildkite/scripts/lib/select-python.sh
select_python
# Detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)" 2>/dev/null || echo 0)
if [[ "$has_new_python" -eq 0 ]]; then
# Use new python from docker
# Use --user to ensure files are created with correct ownership (not root)
docker pull python:3-slim
PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
fi
# Set up auditwheel-in-a-container for the manylinux retagging step.
# Distinct from select_python: ``manylinux.sh`` deliberately pins both
# the Python and auditwheel versions (the script reads auditwheel
# internals) and so always runs in a known-good container regardless
# of what's on the agent.
# shellcheck source=lib/manylinux.sh
source .buildkite/scripts/lib/manylinux.sh
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== Part 1: Collect and prepare wheels ========
@@ -61,18 +63,11 @@ if [ "$WHEEL_COUNT" -eq 0 ]; then
exit 1
fi
# Detect the appropriate manylinux platform tag for any wheel that still
# carries the generic ``linux_<arch>`` tag, and rename it in place. We use
# auditwheel via ``apply_manylinux_tag`` (see lib/manylinux.sh) rather than
# a hard-coded ``manylinux_2_35`` string so that the label tracks the actual
# glibc symbol versions used by the binaries (and stays correct if the
# rocm_base image is rebased).
#
# The ``linux``/``manylinux`` filter below skips both pre-tagged wheels
# (e.g. upstream torch) and pure-Python ``-any.whl`` wheels.
# Rename linux to manylinux in wheel filenames
for wheel in all-rocm-wheels/*.whl; do
if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
new_wheel="$(apply_manylinux_tag "$wheel")"
new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
mv -- "$wheel" "$new_wheel"
echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
fi
done
+31 -31
View File
@@ -388,18 +388,18 @@ 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
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
#---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------#
@@ -1168,13 +1168,13 @@ steps:
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- tests/distributed/test_context_parallel.py
- examples/features/data_parallel/data_parallel_offline.py
- examples/offline_inference/data_parallel.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- pytest -v -s tests/distributed/test_context_parallel.py
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
- label: Distributed Tests (4xA100-4xMI300) # TBD
timeout_in_minutes: 180
@@ -1203,7 +1203,7 @@ steps:
- tests/distributed/test_torchrun_example.py
- tests/distributed/test_torchrun_example_moe.py
- examples/rl/
- tests/examples/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.py
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
@@ -1213,7 +1213,7 @@ steps:
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
- python3 ../examples/features/data_parallel/data_parallel_offline.py --enforce-eager
- python3 ../examples/offline_inference/data_parallel.py --enforce-eager
# rlhf examples
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_ipc.py
@@ -1266,7 +1266,7 @@ steps:
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- examples/features/torchrun/torchrun_dp_example_offline.py
- examples/offline_inference/torchrun_dp_example.py
- vllm/config/parallel.py
- vllm/distributed/
- vllm/v1/engine/llm_engine.py
@@ -1275,7 +1275,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
@@ -1647,18 +1647,18 @@ 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
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
@@ -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 --------------------------------------------------------#
@@ -2302,7 +2302,7 @@ steps:
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
@@ -2713,7 +2713,7 @@ steps:
- vllm/v1/attention/selector.py
- tests/distributed/test_context_parallel.py
- tests/v1/distributed/test_dbo.py
- examples/features/data_parallel/data_parallel_offline.py
- examples/offline_inference/data_parallel.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
@@ -2930,18 +2930,18 @@ 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
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: V1 attention (H100)
key: v1-attention-h100
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -15,7 +14,6 @@ steps:
- pytest -v -s v1/attention
- label: V1 attention (B200)
key: v1-attention-b200
timeout_in_minutes: 30
device: b200
source_file_dependencies:
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
@@ -13,7 +12,6 @@ steps:
- pytest -v -s benchmarks/
- label: Attention Benchmarks Smoke Test (B200)
key: attention-benchmarks-smoke-test-b200
device: b200
num_gpus: 2
optional: true
-13
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Sequence Parallel Correctness Tests (2 GPUs)
key: sequence-parallel-correctness-tests-2-gpus
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
num_devices: 2
@@ -18,7 +17,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: Sequence Parallel Correctness Tests (2xH100)
key: sequence-parallel-correctness-tests-2xh100
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: h100
@@ -29,7 +27,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: AsyncTP Correctness Tests (2xH100)
key: asynctp-correctness-tests-2xh100
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: h100
@@ -40,7 +37,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: AsyncTP Correctness Tests (B200)
key: asynctp-correctness-tests-b200
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: b200
@@ -51,7 +47,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: Distributed Compile Unit Tests (2xH100)
key: distributed-compile-unit-tests-2xh100
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: h100
@@ -65,7 +60,6 @@ steps:
- pytest -s -v tests/compile/passes/distributed
- label: Fusion and Compile Unit Tests (2xB200)
key: fusion-and-compile-unit-tests-2xb200
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: b200
@@ -95,7 +89,6 @@ steps:
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
- label: Fusion E2E Quick (H100)
key: fusion-e2e-quick-h100
timeout_in_minutes: 15
working_dir: "/vllm-workspace/"
device: h100
@@ -114,7 +107,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
- label: Fusion E2E Config Sweep (H100)
key: fusion-e2e-config-sweep-h100
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: h100
@@ -134,7 +126,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
- label: Fusion E2E Config Sweep (B200)
key: fusion-e2e-config-sweep-b200
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: b200
@@ -148,7 +139,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"
- label: Fusion E2E TP2 Quick (H100)
key: fusion-e2e-tp2-quick-h100
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: h100
@@ -166,7 +156,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
key: fusion-e2e-tp2-ar-rms-config-sweep-h100
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
device: h100
@@ -186,7 +175,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
- label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
key: fusion-e2e-tp2-asynctp-config-sweep-h100
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
device: h100
@@ -206,7 +194,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
- label: Fusion E2E TP2 (B200)
key: fusion-e2e-tp2-b200
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: b200
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Platform Tests (CUDA)
key: platform-tests-cuda
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s cuda/test_platform_no_cuda_init.py
- label: Cudagraph
key: cudagraph
timeout_in_minutes: 20
source_file_dependencies:
- tests/v1/cudagraph
-8
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Distributed NixlConnector PD accuracy (4 GPUs)
key: distributed-nixlconnector-pd-accuracy-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -14,7 +13,6 @@ steps:
- 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)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -26,7 +24,6 @@ steps:
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -38,7 +35,6 @@ steps:
- 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)
key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -50,7 +46,6 @@ steps:
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -62,7 +57,6 @@ steps:
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -77,7 +71,6 @@ steps:
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
key: nixlconnector-pd-spec-decode-acceptance-2-gpus
timeout_in_minutes: 30
device: a100
working_dir: "/vllm-workspace/tests"
@@ -91,7 +84,6 @@ steps:
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
+8 -26
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Distributed Comm Ops
key: distributed-comm-ops
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -17,7 +16,6 @@ steps:
- pytest -v -s distributed/test_shm_storage.py
- label: Distributed DP Tests (2 GPUs)
key: distributed-dp-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -39,7 +37,6 @@ steps:
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
- label: Distributed Compile + RPC Tests (2 GPUs)
key: distributed-compile-rpc-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -62,7 +59,6 @@ steps:
- pytest -v -s ./compile/test_wrapper.py
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
key: distributed-torchrun-shutdown-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -85,7 +81,6 @@ steps:
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: Distributed Torchrun + Examples (4 GPUs)
key: distributed-torchrun-examples-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace"
num_devices: 4
@@ -93,8 +88,9 @@ steps:
- vllm/distributed/
- tests/distributed/test_torchrun_example.py
- tests/distributed/test_torchrun_example_moe.py
- examples/offline_inference/rlhf_colocate.py
- examples/rl/
- tests/examples/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -111,13 +107,12 @@ steps:
# test with torchrun tp=2 and dp=2 with ep
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with internal dp
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
- python3 examples/offline_inference/data_parallel.py --enforce-eager
# rlhf examples
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
- label: Distributed DP Tests (4 GPUs)
key: distributed-dp-tests-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -138,7 +133,6 @@ steps:
- pytest -v -s distributed/test_utils.py
- label: Distributed Compile + Comm (4 GPUs)
key: distributed-compile-comm-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -160,28 +154,24 @@ steps:
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
- label: Distributed Tests (8 GPUs)(H100)
key: distributed-tests-8-gpus-h100
timeout_in_minutes: 10
device: h100
num_devices: 8
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- examples/features/torchrun/torchrun_dp_example_offline.py
- examples/offline_inference/torchrun_dp_example.py
- vllm/config/parallel.py
- vllm/distributed/
- vllm/v1/engine/llm_engine.py
- vllm/v1/executor/uniproc_executor.py
- vllm/v1/worker/gpu_worker.py
- tests/distributed/test_mnnvl_alltoall.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
# test with torchrun tp=2 and dp=4 with ep
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- label: Distributed Tests (4 GPUs)(A100)
key: distributed-tests-4-gpus-a100
device: a100
optional: true
num_devices: 4
@@ -196,7 +186,6 @@ steps:
- pytest -v -s -x lora/test_mixtral.py
- label: Distributed Tests (2 GPUs)(H100)
key: distributed-tests-2-gpus-h100
timeout_in_minutes: 15
device: h100
optional: true
@@ -205,13 +194,12 @@ steps:
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- label: Distributed Tests (2 GPUs)(B200)
key: distributed-tests-2-gpus-b200
device: b200
optional: true
working_dir: "/vllm-workspace/"
@@ -220,12 +208,8 @@ steps:
- pytest -v -s tests/distributed/test_context_parallel.py
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
- pytest -v -s tests/v1/distributed/test_dbo.py
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
- label: 2 Node Test (4 GPUs)
key: 2-node-test-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -238,12 +222,11 @@ steps:
- vllm/executor/
- vllm/model_executor/models/
- tests/distributed/
- tests/examples/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.py
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/features/data_parallel/data_parallel_offline.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/features/data_parallel/data_parallel_offline.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"
- ./.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: Pipeline + Context Parallelism (4 GPUs)
key: pipeline-context-parallelism-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -258,7 +241,6 @@ steps:
- pytest -v -s distributed/test_pipeline_parallel.py
- label: RayExecutorV2 (4 GPUs)
key: rayexecutorv2-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
-16
View File
@@ -1,16 +0,0 @@
group: Docker
depends_on:
- image-build-cpu
steps:
- label: Docker Build Metadata
timeout_in_minutes: 10
device: cpu-small
source_file_dependencies:
- .buildkite/release-pipeline.yaml
- .buildkite/scripts/docker-build-metadata-args.sh
- docker/Dockerfile
- docker/Dockerfile.cpu
- docker/docker-bake.hcl
- tests/tools/test_docker_build_metadata_args.py
commands:
- pytest -v -s tools/test_docker_build_metadata_args.py
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: DeepSeek V2-Lite Accuracy
key: deepseek-v2-lite-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -13,7 +12,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
- label: Qwen3-30B-A3B-FP8-block Accuracy
key: qwen3-30b-a3b-fp8-block-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -23,7 +21,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
key: qwen3-30b-a3b-fp8-block-accuracy-b200
timeout_in_minutes: 60
device: b200
optional: true
@@ -33,7 +30,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -43,7 +39,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
key: deepseek-v2-lite-prefetch-offload-accuracy-h100
timeout_in_minutes: 60
device: h100
optional: true
-7
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Engine
key: engine
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
@@ -17,7 +16,6 @@ steps:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
- label: Engine (1 GPU)
key: engine-1-gpu
timeout_in_minutes: 30
source_file_dependencies:
- vllm/v1/engine/
@@ -27,7 +25,6 @@ steps:
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -37,7 +34,6 @@ steps:
- pytest -v -s v1/e2e/general/test_async_scheduling.py
- label: e2e Core (1 GPU)
key: e2e-core-1-gpu
timeout_in_minutes: 30
source_file_dependencies:
- vllm/v1/
@@ -46,7 +42,6 @@ steps:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
optional: true
num_devices: 2
@@ -63,7 +58,6 @@ steps:
- image-build-amd
- label: V1 e2e (4 GPUs)
key: v1-e2e-4-gpus
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
optional: true
num_devices: 4
@@ -80,7 +74,6 @@ steps:
- image-build-amd
- label: V1 e2e (4xH100)
key: v1-e2e-4xh100
timeout_in_minutes: 60
device: h100
num_devices: 4
+1 -10
View File
@@ -2,8 +2,7 @@ group: Entrypoints
depends_on:
- image-build
steps:
- label: Entrypoints Unit Tests
key: entrypoints-unit-tests
- label: Entrypoints Unit Tests
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -28,7 +26,6 @@ steps:
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
- label: Entrypoints Integration (API Server openai - Part 1)
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -46,7 +43,6 @@ steps:
- label: Entrypoints Integration (API Server openai - Part 2)
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -64,7 +60,6 @@ steps:
- image-build-amd
- label: Entrypoints Integration (API Server openai - Part 3)
key: entrypoints-integration-api-server-openai-part-3
timeout_in_minutes: 50
device: h200_18gb
working_dir: "/vllm-workspace/tests"
@@ -77,7 +72,6 @@ steps:
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (API Server 2)
key: entrypoints-integration-api-server-2
timeout_in_minutes: 130
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -92,7 +86,6 @@ steps:
- pytest -v -s tool_use
- label: Entrypoints Integration (Pooling)
key: entrypoints-integration-pooling
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -103,7 +96,6 @@ steps:
- pytest -v -s entrypoints/pooling
- label: Entrypoints Integration (Responses API)
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -113,7 +105,6 @@ steps:
- pytest -v -s entrypoints/openai/responses
- label: OpenAI API Correctness
key: openai-api-correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: EPLB Algorithm
key: eplb-algorithm
timeout_in_minutes: 15
device: h200_18gb
working_dir: "/vllm-workspace/tests"
@@ -16,7 +15,6 @@ steps:
- pytest -v -s distributed/test_eplb_utils.py
- label: EPLB Execution # 17min
key: eplb-execution
timeout_in_minutes: 27
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -28,7 +26,6 @@ steps:
- pytest -v -s distributed/test_eplb_spec_decode.py
- label: Elastic EP Scaling Test
key: elastic-ep-scaling-test
timeout_in_minutes: 20
device: h100
working_dir: "/vllm-workspace/tests"
-27
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: vLLM IR Tests
key: vllm-ir-tests
timeout_in_minutes: 10
device: h200_18gb
working_dir: "/vllm-workspace/"
@@ -15,7 +14,6 @@ steps:
- pytest -v -s tests/kernels/ir
- label: Kernels Core Operation Test
key: kernels-core-operation-test
timeout_in_minutes: 75
source_file_dependencies:
- csrc/
@@ -25,7 +23,6 @@ steps:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
key: kernels-minimax-reduce-rms-test-2-gpus
timeout_in_minutes: 15
num_devices: 2
device: h100
@@ -39,7 +36,6 @@ steps:
- pytest -v -s kernels/core/test_minimax_reduce_rms.py
- label: Kernels Attention Test %N
key: kernels-attention-test
timeout_in_minutes: 35
source_file_dependencies:
- csrc/attention/
@@ -53,7 +49,6 @@ steps:
parallelism: 2
- label: Kernels Quantization Test %N
key: kernels-quantization-test
timeout_in_minutes: 90
source_file_dependencies:
- csrc/quantization/
@@ -64,7 +59,6 @@ steps:
parallelism: 2
- label: Kernels MoE Test %N
key: kernels-moe-test
timeout_in_minutes: 25
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
@@ -80,7 +74,6 @@ steps:
parallelism: 5
- label: Kernels Mamba Test
key: kernels-mamba-test
timeout_in_minutes: 45
source_file_dependencies:
- csrc/mamba/
@@ -89,18 +82,7 @@ steps:
commands:
- pytest -v -s kernels/mamba
- label: Kernels KDA Test
timeout_in_minutes: 20
source_file_dependencies:
- vllm/model_executor/layers/fla/ops/kda.py
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
- vllm/model_executor/layers/fla/ops/l2norm.py
- tests/kernels/test_kda.py
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels DeepGEMM Test (H100)
key: kernels-deepgemm-test-h100
timeout_in_minutes: 45
device: h100
num_devices: 1
@@ -113,16 +95,13 @@ 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)
key: kernels-b200
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: b200
@@ -171,7 +150,6 @@ steps:
- pytest -v -s tests/models/quantization/test_nvfp4.py
- label: Kernels Helion Test
key: kernels-helion-test
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -183,7 +161,6 @@ steps:
- label: Kernels FP8 MoE Test (1 H100)
key: kernels-fp8-moe-test-1-h100
timeout_in_minutes: 90
device: h100
num_devices: 1
@@ -200,7 +177,6 @@ steps:
- pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
- label: Kernels FP8 MoE Test (2 H100s)
key: kernels-fp8-moe-test-2-h100s
timeout_in_minutes: 90
device: h100
num_devices: 2
@@ -210,7 +186,6 @@ steps:
- pytest -v -s kernels/moe/test_deepep_moe.py
- label: Kernels Fp4 MoE Test (B200)
key: kernels-fp4-moe-test-b200
timeout_in_minutes: 60
device: b200
num_devices: 1
@@ -223,7 +198,6 @@ steps:
- label: Kernels FusedMoE Layer Test (2 H100s)
key: kernels-fusedmoe-layer-test-2-h100s
timeout_in_minutes: 90
device: h100
num_devices: 2
@@ -240,7 +214,6 @@ steps:
- label: Kernels FusedMoE Layer Test (2 B200s)
key: kernels-fusedmoe-layer-test-2-b200s
timeout_in_minutes: 90
device: b200
num_devices: 2
-11
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: LM Eval Small Models
key: lm-eval-small-models
timeout_in_minutes: 75
source_file_dependencies:
- csrc/
@@ -25,7 +24,6 @@ steps:
# - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: LM Eval Large Models (4 GPUs)(H100)
key: lm-eval-large-models-4-gpus-h100
device: h100
optional: true
num_devices: 4
@@ -38,7 +36,6 @@ steps:
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
- label: LM Eval Small Models (B200)
key: lm-eval-small-models-b200
timeout_in_minutes: 120
device: b200
optional: true
@@ -49,7 +46,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
- label: LM Eval Qwen3.5 Models (B200)
key: lm-eval-qwen3-5-models-b200
timeout_in_minutes: 120
device: b200
optional: true
@@ -66,7 +62,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
- label: LM Eval Large Models (H200)
key: lm-eval-large-models-h200
timeout_in_minutes: 60
device: h200
optional: true
@@ -75,7 +70,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
device: h100
optional: true
num_devices: 2
@@ -83,7 +77,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
- label: MoE Refactor Integration Test (B200 - TEMPORARY)
key: moe-refactor-integration-test-b200-temporary
device: b200
optional: true
num_devices: 2
@@ -91,7 +84,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
key: moe-refactor-integration-test-b200-dp-temporary
device: b200
optional: true
num_devices: 2
@@ -100,7 +92,6 @@ steps:
- label: LM Eval TurboQuant KV Cache
key: lm-eval-turboquant-kv-cache
timeout_in_minutes: 75
source_file_dependencies:
- vllm/model_executor/layers/quantization/turboquant/
@@ -111,7 +102,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
- label: GPQA Eval (GPT-OSS) (H100)
key: gpqa-eval-gpt-oss-h100
timeout_in_minutes: 120
device: h100
optional: true
@@ -125,7 +115,6 @@ steps:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
- label: GPQA Eval (GPT-OSS) (B200)
key: gpqa-eval-gpt-oss-b200
timeout_in_minutes: 120
device: b200
optional: true
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: LoRA %N
key: lora
timeout_in_minutes: 30
source_file_dependencies:
- vllm/lora
@@ -14,7 +13,6 @@ steps:
- label: LoRA TP (Distributed)
key: lora-tp-distributed
timeout_in_minutes: 30
num_devices: 4
source_file_dependencies:
+8 -21
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: V1 Spec Decode
key: v1-spec-decode
timeout_in_minutes: 30
source_file_dependencies:
- vllm/
@@ -19,7 +18,6 @@ steps:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -43,7 +41,6 @@ steps:
- image-build-amd
- label: V1 Core + KV + Metrics
key: v1-core-kv-metrics
timeout_in_minutes: 30
source_file_dependencies:
- vllm/
@@ -74,7 +71,6 @@ steps:
- image-build-amd
- label: V1 Others (CPU)
key: v1-others-cpu
depends_on:
- image-build-cpu
source_file_dependencies:
@@ -90,7 +86,6 @@ steps:
- pytest -v -s -m 'cpu_test' v1/metrics
- label: Regression
key: regression
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
@@ -102,7 +97,6 @@ steps:
working_dir: "/vllm-workspace/tests" # optional
- label: Examples
key: examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
@@ -119,22 +113,21 @@ 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
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- label: Metrics, Tracing (2 GPUs)
key: metrics-tracing-2-gpus
timeout_in_minutes: 20
num_devices: 2
source_file_dependencies:
@@ -149,7 +142,6 @@ steps:
- pytest -v -s v1/tracing
- label: Python-only Installation
key: python-only-installation
depends_on: ~
timeout_in_minutes: 20
source_file_dependencies:
@@ -159,7 +151,6 @@ steps:
- bash standalone_tests/python_only_compile.sh
- label: Async Engine, Inputs, Utils, Worker
key: async-engine-inputs-utils-worker
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
@@ -172,8 +163,7 @@ steps:
- pytest -v -s utils_
- label: Async Engine, Inputs, Utils, Worker, Config (CPU)
key: async-engine-inputs-utils-worker-config-cpu
depends_on:
depends_on:
- image-build-cpu
timeout_in_minutes: 30
source_file_dependencies:
@@ -206,7 +196,6 @@ steps:
- pytest -v -s config
- label: Batch Invariance (H100)
key: batch-invariance-h100
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -222,7 +211,6 @@ steps:
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- label: Batch Invariance (B200)
key: batch-invariance-b200
timeout_in_minutes: 30
device: b200
source_file_dependencies:
@@ -239,7 +227,6 @@ steps:
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- label: Acceptance Length Test (Large Models) # optional
key: acceptance-length-test-large-models
timeout_in_minutes: 25
gpu: h100
optional: true
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Model Executor
key: model-executor
timeout_in_minutes: 35
source_file_dependencies:
- vllm/engine/arg_utils.py
+8 -14
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Model Runner V2 Core Tests
key: model-runner-v2-core-tests
timeout_in_minutes: 45
source_file_dependencies:
- vllm/v1/worker/gpu/
@@ -26,16 +25,14 @@ steps:
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
- label: Model Runner V2 Examples
key: model-runner-v2-examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/core/sched/
- vllm/v1/worker/gpu_worker.py
- examples/offline_inference/
- examples/basic/offline_inference/
- examples/generate/multimodal/
- examples/features/
- examples/pooling/embed/vision_embedding_offline.py
- examples/others/tensorize_vllm_model.py
commands:
@@ -47,22 +44,21 @@ 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
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- label: Model Runner V2 Distributed (2 GPUs)
key: model-runner-v2-distributed-2-gpus
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -83,7 +79,6 @@ steps:
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
key: model-runner-v2-pipeline-parallelism-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -99,7 +94,6 @@ steps:
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- label: Model Runner V2 Spec Decode
key: model-runner-v2-spec-decode
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
+5 -11
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Basic Models Tests (Initialization)
key: basic-models-tests-initialization
timeout_in_minutes: 45
torch_nightly: true
source_file_dependencies:
@@ -17,7 +16,6 @@ steps:
torch_nightly: {}
- label: Basic Models Tests (Extra Initialization) %N
key: basic-models-tests-extra-initialization
timeout_in_minutes: 45
source_file_dependencies:
- vllm/model_executor/models/
@@ -33,7 +31,6 @@ steps:
torch_nightly: {}
- label: Basic Models Tests (Other)
key: basic-models-tests-other
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -50,7 +47,6 @@ steps:
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
- image-build-cpu
timeout_in_minutes: 10
@@ -63,7 +59,6 @@ steps:
- pytest -v -s models/test_utils.py models/test_vision.py
- label: Transformers Nightly Models
key: transformers-nightly-models
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
@@ -74,12 +69,11 @@ 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
key: transformers-backward-compatibility-models-test
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
@@ -89,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
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Distributed Model Tests (2 GPUs)
key: distributed-model-tests-2-gpus
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Language Models Tests (Standard)
key: language-models-tests-standard
timeout_in_minutes: 25
source_file_dependencies:
- vllm/
@@ -16,7 +15,6 @@ steps:
torch_nightly: {}
- label: Language Models Tests (Extra Standard) %N
key: language-models-tests-extra-standard
timeout_in_minutes: 45
source_file_dependencies:
- vllm/model_executor/models/
@@ -33,7 +31,6 @@ steps:
torch_nightly: {}
- label: Language Models Tests (Hybrid) %N
key: language-models-tests-hybrid
timeout_in_minutes: 75
source_file_dependencies:
- vllm/
@@ -50,7 +47,6 @@ steps:
torch_nightly: {}
- label: Language Models Test (Extended Generation) # 80min
key: language-models-test-extended-generation
timeout_in_minutes: 110
optional: true
source_file_dependencies:
@@ -73,7 +69,6 @@ steps:
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
- label: Language Models Test (PPL)
key: language-models-test-ppl
timeout_in_minutes: 110
device: h200_18gb
optional: true
@@ -84,7 +79,6 @@ steps:
- pytest -v -s models/language/generation_ppl_test
- label: Language Models Test (Extended Pooling) # 36min
key: language-models-test-extended-pooling
timeout_in_minutes: 50
optional: true
source_file_dependencies:
@@ -99,7 +93,6 @@ steps:
- image-build-amd
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 110
device: h200_18gb
optional: true
+1 -12
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
@@ -20,7 +19,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
@@ -38,7 +36,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -54,7 +51,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 4: other + whisper"
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -71,8 +67,7 @@ steps:
- image-build-amd
- label: Multi-Modal Processor (CPU)
key: multi-modal-processor-cpu
depends_on:
depends_on:
- image-build-cpu
timeout_in_minutes: 60
source_file_dependencies:
@@ -85,7 +80,6 @@ steps:
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 60
device: h200_18gb
source_file_dependencies:
@@ -97,7 +91,6 @@ steps:
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
key: multi-modal-accuracy-eval-small-models
timeout_in_minutes: 70
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
@@ -108,7 +101,6 @@ steps:
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
- label: Multi-Modal Models (Extended Generation 1)
key: multi-modal-models-extended-generation-1
optional: true
source_file_dependencies:
- vllm/
@@ -125,7 +117,6 @@ steps:
- image-build-amd
- label: Multi-Modal Models (Extended Generation 2)
key: multi-modal-models-extended-generation-2
optional: true
source_file_dependencies:
- vllm/
@@ -135,7 +126,6 @@ steps:
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
- label: Multi-Modal Models (Extended Generation 3)
key: multi-modal-models-extended-generation-3
optional: true
source_file_dependencies:
- vllm/
@@ -145,7 +135,6 @@ steps:
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: Multi-Modal Models (Extended Pooling)
key: multi-modal-models-extended-pooling
optional: true
device: h200_18gb
source_file_dependencies:
-1
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Plugin Tests (2 GPUs)
key: plugin-tests-2-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 2
-6
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: PyTorch Compilation Unit Tests
key: pytorch-compilation-unit-tests
timeout_in_minutes: 10
source_file_dependencies:
- vllm/
@@ -19,7 +18,6 @@ steps:
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Compilation Unit Tests (H100)
key: pytorch-compilation-unit-tests-h100
timeout_in_minutes: 30
device: h100
num_devices: 1
@@ -30,7 +28,6 @@ steps:
- "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Compilation Passes Unit Tests
key: pytorch-compilation-passes-unit-tests
timeout_in_minutes: 20
source_file_dependencies:
- vllm/
@@ -39,7 +36,6 @@ steps:
- pytest -s -v compile/passes --ignore compile/passes/distributed
- label: PyTorch Fullgraph Smoke Test
key: pytorch-fullgraph-smoke-test
timeout_in_minutes: 35
source_file_dependencies:
- vllm/
@@ -52,7 +48,6 @@ steps:
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Fullgraph
key: pytorch-fullgraph
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -63,7 +58,6 @@ steps:
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
- label: Pytorch Nightly Dependency Override Check # 2min
key: pytorch-nightly-dependency-override-check
# if this test fails, it means the nightly torch version is not compatible with some
# of the dependencies. Please check the error message and add the package to whitelist
# in /vllm/tools/pre_commit/generate_nightly_torch_test.py
-3
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Quantization
key: quantization
timeout_in_minutes: 90
source_file_dependencies:
- csrc/
@@ -22,7 +21,6 @@ steps:
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
- label: Quantized MoE Test (B200)
key: quantized-moe-test-b200
timeout_in_minutes: 60
working_dir: "/vllm-workspace/"
device: b200
@@ -40,7 +38,6 @@ steps:
- pytest -s -v tests/quantization/test_blackwell_moe.py
- label: Quantized Models Test
key: quantized-models-test
timeout_in_minutes: 60
source_file_dependencies:
- vllm/model_executor/layers/quantization
-1
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Ray Dependency Compatibility Check
key: ray-dependency-compatibility-check
# Informational only — does not block the pipeline.
# If this fails, it means the PR introduces a dependency that
# conflicts with Ray's dependency constraints.
+1 -4
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Samplers Test
key: samplers-test
timeout_in_minutes: 75
source_file_dependencies:
- vllm/model_executor/layers
@@ -11,9 +10,7 @@ steps:
- tests/samplers
- tests/conftest.py
commands:
# VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
# values explicitly to still cover the PyTorch-native (Triton) path.
- VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
- pytest -v -s samplers
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
mirror:
amd:
-21
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Spec Decode Eagle
key: spec-decode-eagle
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Eagle Nightly B200
key: spec-decode-eagle-nightly-b200
timeout_in_minutes: 30
device: b200
optional: true
@@ -26,7 +24,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Speculators + MTP
key: spec-decode-speculators-mtp
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -38,7 +35,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Speculators + MTP Nightly B200
key: spec-decode-speculators-mtp-nightly-b200
timeout_in_minutes: 30
device: b200
optional: true
@@ -51,7 +47,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Ngram + Suffix
key: spec-decode-ngram-suffix
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -62,7 +57,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- label: Spec Decode Draft Model
key: spec-decode-draft-model
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -73,7 +67,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Draft Model Nightly B200
key: spec-decode-draft-model-nightly-b200
timeout_in_minutes: 30
device: b200
optional: true
@@ -85,7 +78,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: DFlash Speculators Correctness
key: dflash-speculators-correctness
timeout_in_minutes: 30
device: h100
optional: true
@@ -97,16 +89,3 @@ steps:
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/qwen3_5_mtp.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Weight Loading Multiple GPU # 33min
key: weight-loading-multiple-gpu
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
num_devices: 2
+8 -3
View File
@@ -308,7 +308,8 @@ pull_request_rules:
- files=benchmarks/benchmark_serving_structured_output.py
- files=benchmarks/run_structured_output_benchmark.sh
- files=docs/features/structured_outputs.md
- files=^examples/features/structured_outputs/
- files=examples/offline_inference/structured_outputs.py
- files=examples/online_serving/structured_outputs/structured_outputs.py
- files~=^tests/v1/structured_output/
- files=tests/entrypoints/llm/test_struct_output_generate.py
- files~=^vllm/v1/structured_output/
@@ -324,7 +325,7 @@ pull_request_rules:
- or:
- files~=^vllm/v1/spec_decode/
- files~=^tests/v1/spec_decode/
- files=^examples/features/speculative_decoding/
- files~=^examples/.*(spec_decode|mlpspeculator|eagle|speculation).*\.py
- files~=^vllm/model_executor/models/.*eagle.*\.py
- files=vllm/model_executor/models/mlp_speculator.py
- files~=^vllm/transformers_utils/configs/(eagle|medusa|mlp_speculator)\.py
@@ -388,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:
+11 -2
View File
@@ -16,7 +16,11 @@ permissions:
jobs:
pre-run-check:
if: github.event_name == 'pull_request'
if: >-
github.event_name == 'pull_request' &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified')
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
@@ -45,7 +49,12 @@ jobs:
pre-commit:
needs: pre-run-check
if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
if: >-
always() &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified') &&
(needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
-1
View File
@@ -237,7 +237,6 @@ ep_kernels_workspace/
# Allow tracked library source folders under submodules (e.g., benchmarks/lib)
!vllm/benchmarks/lib/
!.buildkite/scripts/lib/
# Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
vllm/grpc/vllm_engine_pb2.py
+21
View File
@@ -564,6 +564,27 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
"in CUDA target architectures.")
endif()
# DeepSeek V4 indexer top-k. Needs thread-block clusters + TMA + PDL, so
# builds for Hopper (sm_90a) and Blackwell datacenter (sm_100/sm_103). Not
# supported on sm_120 (consumer Blackwell, no clusters). Requires CUDA >=
# 12.4 for the cuda::ptx mbarrier wrappers. Ported from sglang.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.4 AND DSV4_TOPK_ARCHS)
set(DSV4_TOPK_SRC "csrc/deepseek_v4/fast_topk_v2.cu")
set_gencode_flags_for_srcs(
SRCS "${DSV4_TOPK_SRC}"
CUDA_ARCHS "${DSV4_TOPK_ARCHS}")
list(APPEND VLLM_EXT_SRC ${DSV4_TOPK_SRC})
message(STATUS "Building deepseek_v4 fast_topk_v2 for archs: ${DSV4_TOPK_ARCHS}")
else()
message(STATUS "Not building deepseek_v4 fast_topk_v2 (needs CUDA >= 12.4 "
"and a compatible Hopper+ arch).")
endif()
#
# Machete kernels
@@ -0,0 +1,183 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Microbench: fast_topk_v2 vs persistent_topk for k in {512, 1024}.
Both ops select the top-k entries per row of a `[B, L]` float32 score
tensor. vLLM's `persistent_topk` is the existing path used by the indexer;
`fast_topk_v2` is the sm_90+ port from sglang that adds Hopper thread-block
clusters.
V4-Flash uses `index_topk = 512`; V4-Pro uses `index_topk = 1024`. We bench
both Ks at the realistic shape regimes (small-B, L up to 256K compressed).
Timing uses **CUDA graph replay** to amortize launch overhead (~3-5 µs on
Blackwell). We capture N invocations of the same kernel, replay the graph
many times, divide.
Run::
.venv/bin/python benchmarks/kernels/benchmark_fast_topk_v2.py
"""
from __future__ import annotations
import argparse
import statistics
import sys
import torch
import vllm._C # noqa: F401 ensures schemas are registered
from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import (
fast_topk_v2_raw,
plan_topk_v2,
workspace_ints_per_batch,
)
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 # bytes; matches sparse_attn_indexer.py
def _capture_graph(callable_fn, *, calls_per_graph: int) -> torch.cuda.CUDAGraph:
for _ in range(3):
callable_fn()
torch.cuda.synchronize()
g = torch.cuda.CUDAGraph()
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
with torch.cuda.graph(g, stream=s):
for _ in range(calls_per_graph):
callable_fn()
torch.cuda.current_stream().wait_stream(s)
return g
def time_graph_us(graph: torch.cuda.CUDAGraph, *, calls_per_graph: int,
warmup: int = 5, replays: int = 30) -> float:
for _ in range(warmup):
graph.replay()
torch.cuda.synchronize()
samples = []
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
for _ in range(replays):
start.record()
graph.replay()
end.record()
end.synchronize()
samples.append(start.elapsed_time(end) * 1000.0 / calls_per_graph)
return statistics.median(samples)
def make_inputs(batch_size: int, seq_len: int, *, seed: int = 0):
device = torch.device("cuda")
g = torch.Generator(device=device).manual_seed(seed)
L = (seq_len + 3) & ~3
scores = torch.randn(batch_size, L, generator=g, dtype=torch.float32,
device=device)
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32,
device=device)
return scores, seq_lens, L
def bench_persistent_topk(scores, seq_lens, k, *, calls_per_graph: int) -> float:
B = scores.shape[0]
output = scores.new_empty((B, k), dtype=torch.int32)
workspace = scores.new_empty((RADIX_TOPK_WORKSPACE_SIZE,), dtype=torch.uint8)
max_seq_len = scores.shape[1]
def run():
torch.ops._C.persistent_topk(
scores, seq_lens, output, workspace, k, max_seq_len)
graph = _capture_graph(run, calls_per_graph=calls_per_graph)
return time_graph_us(graph, calls_per_graph=calls_per_graph)
def bench_fast_topk_v2(scores, seq_lens, k, *,
calls_per_graph: int) -> float:
B = scores.shape[0]
metadata = plan_topk_v2(seq_lens)
workspace = scores.new_empty((B, workspace_ints_per_batch()),
dtype=torch.int32)
topk_indices = scores.new_empty((B, k), dtype=torch.int32)
def run():
fast_topk_v2_raw(scores, seq_lens, topk=k,
metadata=metadata, workspace=workspace,
topk_indices=topk_indices)
graph = _capture_graph(run, calls_per_graph=calls_per_graph)
return time_graph_us(graph, calls_per_graph=calls_per_graph)
def fmt(us: float) -> str:
return f"{us:8.2f}"
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--batch-sizes", type=int, nargs="+",
default=[1, 4, 16, 32, 64, 128, 256])
parser.add_argument("--seq-lens", type=int, nargs="+",
default=[1024, 4096, 16384, 32768, 65536, 131072])
parser.add_argument("--ks", type=int, nargs="+",
default=[512, 1024])
parser.add_argument("--calls-per-graph", type=int, default=64)
parser.add_argument("--replays", type=int, default=30)
args = parser.parse_args()
if not torch.cuda.is_available():
print("CUDA is required for this benchmark.", file=sys.stderr)
sys.exit(1)
print(f"GPU: {torch.cuda.get_device_name(0)} "
f"(SM {torch.cuda.get_device_capability(0)})")
print(f"calls_per_graph={args.calls_per_graph}, replays={args.replays}")
print("Per-call medians via CUDA graph replay (host launch overhead "
"amortized).\n")
for k in args.ks:
print(f"=== k = {k} ===")
print(f"{'B':>4} {'L':>7} | {'persistent_topk':>17} | "
f"{'fast_topk_v2':>14} | {'speedup':>8} | {'path':<14}")
print("-" * 80)
for B in args.batch_sizes:
for L in args.seq_lens:
# Skip seq_lens beyond persistent_topk's k-dependent useful
# range. Both kernels handle up to 256K with k=1024.
try:
scores, seq_lens, _ = make_inputs(B, L, seed=B * L * k)
p_us = bench_persistent_topk(
scores, seq_lens, k,
calls_per_graph=args.calls_per_graph)
f_us = bench_fast_topk_v2(
scores, seq_lens, k,
calls_per_graph=args.calls_per_graph)
speedup = p_us / f_us if f_us > 0 else float("inf")
if L <= k:
path = "trivial"
elif L <= 4 * 4 * 1024:
path = "register-1p"
elif L <= 32768:
path = "register-2p"
elif B <= 15:
path = "cluster-fused"
else:
path = "cluster-2stg"
print(
f"{B:>4} {L:>7} | "
f"{fmt(p_us):>14} us | "
f"{fmt(f_us):>11} us | "
f"{speedup:>5.2f}x | {path}"
)
except RuntimeError as e:
print(f"{B:>4} {L:>7} | ERROR: {e}")
print()
if __name__ == "__main__":
main()
@@ -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,
)
@@ -217,7 +217,6 @@ async def send_request(
min_tokens: int | None = None,
max_tokens: int | None = None,
timeout_sec: int = 120,
conversation_id: str | None = None,
) -> ServerResponse:
payload = {
"model": model,
@@ -226,9 +225,6 @@ async def send_request(
"temperature": 0.0,
}
if conversation_id is not None:
payload["conversation_id"] = conversation_id
if stream:
payload["stream"] = True
payload["stream_options"] = {"include_usage": False}
@@ -423,7 +419,6 @@ async def send_turn(
min_tokens,
max_tokens,
req_args.timeout_sec,
conversation_id=conv_id,
)
if response.valid is False:
+25 -82
View File
@@ -11,74 +11,29 @@
namespace vllm {
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y,
const float limit) {
if constexpr (act_first) {
scalar_t gate = x;
scalar_t up = y;
if constexpr (HAS_CLAMP) {
gate = (scalar_t)fminf((float)gate, limit);
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
}
return ACT_FN(gate) * up;
} else {
scalar_t gate = x;
scalar_t up = y;
if constexpr (HAS_CLAMP) {
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
up = (scalar_t)fminf((float)up, limit);
}
return gate * ACT_FN(up);
}
const scalar_t& y) {
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
}
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
const packed_t& y,
const float limit) {
if constexpr (act_first) {
packed_t gate = x;
packed_t up = y;
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fminf(g.x, limit);
g.y = fminf(g.y, limit);
u.x = fmaxf(fminf(u.x, limit), -limit);
u.y = fmaxf(fminf(u.y, limit), -limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(PACKED_ACT_FN(gate), up);
} else {
packed_t gate = x;
packed_t up = y;
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fmaxf(fminf(g.x, limit), -limit);
g.y = fmaxf(fminf(g.y, limit), -limit);
u.x = fminf(u.x, limit);
u.y = fminf(u.y, limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(gate, PACKED_ACT_FN(up));
}
const packed_t& y) {
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
: packed_mul(x, PACKED_ACT_FN(y));
}
// Activation and gating kernel template.
template <typename scalar_t, typename packed_t,
scalar_t (*ACT_FN)(const scalar_t&),
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
bool use_vec, bool use_256b = false>
__global__ void act_and_mul_kernel(
scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d, const float limit) {
const int d) {
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
}
#pragma unroll
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
x.elts[j] =
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
x.elts[j], y.elts[j], limit);
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
x.elts[j], y.elts[j]);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
out_ptr[idx] =
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
}
}
}
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
// Launch activation and gating kernel.
// Use ACT_FIRST (bool) indicating whether to apply the activation function
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
// clamped (max only) and up input is clamped (both sides) before the
// activation function is applied.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
HAS_CLAMP, LIMIT) \
// first.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
auto dtype = input.scalar_type(); \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} \
} else { \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
}
@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, false, 0.0f);
}
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
double limit) {
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, true, (float)limit);
true);
}
void mul_and_silu(torch::Tensor& out, // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
// applies the silu to the latter half of the input.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
false, false, 0.0f);
false);
}
void gelu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
true, false, 0.0f);
true);
}
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true);
}
namespace vllm {
+25 -78
View File
@@ -1,16 +1,5 @@
#include "cpu_attn_dispatch_generated.h"
// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
const std::string& kv_cache_dtype) {
if (kv_cache_dtype == "fp8_e5m2")
return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
return cpu_attention::Fp8KVCacheDataType::kAuto;
}
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
@@ -60,7 +49,7 @@ torch::Tensor get_scheduler_metadata(
input.enable_kv_split = enable_kv_split;
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
input.elem_size = sizeof(scalar_t);
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
@@ -83,9 +72,7 @@ void cpu_attn_reshape_and_cache(
key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
torch::Tensor&
value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
const torch::Tensor& slot_mapping, const std::string& isa,
const double k_scale = 1.0, const double v_scale = 1.0,
const std::string& kv_cache_dtype = "auto") {
const torch::Tensor& slot_mapping, const std::string& isa) {
TORCH_CHECK_EQ(key.dim(), 3);
TORCH_CHECK_EQ(value.dim(), 3);
TORCH_CHECK_EQ(key_cache.dim(), 4);
@@ -93,30 +80,18 @@ void cpu_attn_reshape_and_cache(
TORCH_CHECK_EQ(key.stride(2), 1);
TORCH_CHECK_EQ(value.stride(2), 1);
const int64_t kv_cache_idx =
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
const bool is_fp8 = (kv_cache_idx != 0);
if (is_fp8) {
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
"key_cache must be uint8 for FP8 path");
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
"value_cache must be uint8 for FP8 path");
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
}
const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
const int64_t token_num = key.size(0);
const int64_t head_num = key.size(1);
const int64_t head_dim = key.size(2);
const int64_t key_token_num_stride = key.stride(0);
const int64_t value_token_num_stride = value.stride(0);
const int64_t head_num = value.size(1);
const int64_t key_head_num_stride = key.stride(1);
const int64_t value_head_num_stride = value.stride(1);
const int64_t num_blocks = key_cache.size(0);
const int64_t num_blocks_stride = key_cache.stride(0);
const int64_t cache_head_num_stride = key_cache.stride(1);
const int64_t block_size = key_cache.size(2);
const int64_t block_size_stride = key_cache.stride(2);
const int64_t head_dim = key.size(-1);
cpu_attention::ISA isa_tag = [&]() {
if (isa == "amx") {
@@ -134,24 +109,16 @@ void cpu_attn_reshape_and_cache(
}
}();
if (is_fp8) {
TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
isa_tag == cpu_attention::ISA::VEC,
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
}
VLLM_DISPATCH_FLOATING_TYPES(
key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
using kv_t = typename attn_impl::kv_cache_t;
CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
attn_impl::reshape_and_cache(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
reinterpret_cast<kv_t*>(key_cache.data_ptr()),
reinterpret_cast<kv_t*>(value_cache.data_ptr()),
slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
value.stride(0), head_num, key.stride(1), value.stride(1),
num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
block_size_stride, k_inv, v_inv);
key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
value_token_num_stride, head_num, key_head_num_stride,
value_head_num_stride, num_blocks, num_blocks_stride,
cache_head_num_stride, block_size, block_size_stride);
});
});
}
@@ -170,26 +137,13 @@ void cpu_attention_with_kv_cache(
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, // [num_tokens, max_block_num]
const double softcap, const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, // [num_heads]
const double k_scale = 1.0, const double v_scale = 1.0,
const std::string& kv_cache_dtype = "auto") {
const std::optional<torch::Tensor>& s_aux // [num_heads]
) {
TORCH_CHECK_EQ(query.dim(), 3);
TORCH_CHECK_EQ(query.stride(2), 1);
TORCH_CHECK_EQ(key_cache.dim(), 4);
TORCH_CHECK_EQ(value_cache.dim(), 4);
const int64_t kv_cache_idx =
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
const bool is_fp8 = (kv_cache_idx != 0);
if (is_fp8) {
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
"key_cache must be uint8 for FP8 path");
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
"value_cache must be uint8 for FP8 path");
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
}
cpu_attention::AttentionInput input;
input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
scheduler_metadata.data_ptr());
@@ -211,32 +165,25 @@ void cpu_attention_with_kv_cache(
input.block_table = block_table.data_ptr<int32_t>();
input.alibi_slopes =
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
// For now sink must be bf16
input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
input.scale = scale;
input.causal = causal;
input.sliding_window_left = sliding_window_left;
input.sliding_window_right = sliding_window_right;
if (input.causal) {
// to make boundary calculation easier
input.sliding_window_right = 0;
}
input.softcap = static_cast<float>(softcap);
if (is_fp8) {
input.k_scale_fp8 = static_cast<float>(k_scale);
input.v_scale_fp8 = static_cast<float>(v_scale);
TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
input.metadata->isa == cpu_attention::ISA::VEC,
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
}
float softcap_fp32 = softcap;
input.softcap = softcap_fp32;
VLLM_DISPATCH_FLOATING_TYPES(
query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
CPU_ATTN_DISPATCH(
query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
0);
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
mainloop(&input);
});
CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
mainloop(&input);
});
});
}
+46 -171
View File
@@ -1,7 +1,6 @@
#ifndef CPU_ATTN_AMX_HPP
#define CPU_ATTN_AMX_HPP
#include "cpu_attn_fp8.hpp"
#include "cpu_attn_impl.hpp"
namespace cpu_attention {
@@ -22,10 +21,9 @@ typedef struct __tile_config {
// 2-2-4 pattern, for 16 < m <= 32
// TILE 0, 1: load A matrix, row num should be 16, m - 16
// TILE 2, 3: load B matrix, row num should be 16
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
// m - 16, m - 16
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
template <typename q_buffer_t, typename kv_cache_t>
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
// - 16
template <typename kv_cache_t>
class TileGemm224 {
public:
template <AttentionGemmPhase phase, int32_t k_size>
@@ -44,56 +42,13 @@ class TileGemm224 {
}
};
// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
template <typename kv_cache_t>
FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
.save(dst + r * 32);
} else {
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
.save(dst + r * 32);
}
}
}
// For FP8: dequant src into scratch and return scratch.
// For BF16: return src directly (scratch is unused; the compiler elides it).
template <typename kv_cache_t>
FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
c10::BFloat16* scratch) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
return scratch;
} else {
return reinterpret_cast<const c10::BFloat16*>(src);
}
}
// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
template <typename kv_cache_t>
class TileGemm224<c10::BFloat16, kv_cache_t> {
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
// BF16 path: scratch_elems=1 so the scratch array is eliminated by the
// compiler.
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
template <>
class TileGemm224<c10::BFloat16> {
public:
template <AttentionGemmPhase phase, int32_t k_size>
FORCE_INLINE static void gemm(const int32_t m_size,
c10::BFloat16* __restrict__ a_tile,
kv_cache_t* __restrict__ b_tile,
c10::BFloat16* __restrict__ b_tile,
float* __restrict__ c_tile, const int64_t lda,
const int64_t ldb, const int64_t ldc,
const int32_t block_size,
@@ -101,7 +56,6 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
const bool accum_c) {
const int32_t k_times =
dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
c10::BFloat16* __restrict__ a_tile_0 = a_tile;
c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
const int64_t a_tile_stride = [&]() {
@@ -116,8 +70,8 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
}
}();
kv_cache_t* __restrict__ b_tile_2 = b_tile;
kv_cache_t* __restrict__ b_tile_3 = [&]() {
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
if constexpr (phase == AttentionGemmPhase::QK) {
// k_cache is prepacked
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
@@ -152,16 +106,11 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
_tile_zero(7);
}
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
for (int32_t k = 0; k < k_times; ++k) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
_tile_stream_loadd(2, b_tile_2, b_tile_stride);
_tile_dpbf16ps(4, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
_tile_stream_loadd(3, b_tile_3, b_tile_stride);
_tile_dpbf16ps(5, 0, 3);
_tile_loadd(1, a_tile_1, a_tile_stride);
_tile_dpbf16ps(6, 1, 2);
@@ -205,13 +154,13 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
};
// 1-2-2 pattern, for 0 < m <= 16
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
// be m, m
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
// should be 16
// TILE 6, 7: store results C matrix, row num should be m
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
template <typename q_buffer_t, typename kv_cache_t>
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
// m, m
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
// num should be 16
// TILE 6, 7, (6, 7): store results C matrix, row num should be
// m
template <typename kv_cache_t>
class TileGemm122 {
public:
template <AttentionGemmPhase phase, int32_t k_size>
@@ -230,26 +179,13 @@ class TileGemm122 {
}
};
// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
template <typename kv_cache_t>
class TileGemm122<c10::BFloat16, kv_cache_t> {
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
template <>
class TileGemm122<c10::BFloat16> {
public:
template <AttentionGemmPhase phase, int32_t k_size>
FORCE_INLINE static void gemm(const int32_t m_size,
c10::BFloat16* __restrict__ a_tile,
kv_cache_t* __restrict__ b_tile,
c10::BFloat16* __restrict__ b_tile,
float* __restrict__ c_tile, const int64_t lda,
const int64_t ldb, const int64_t ldc,
const int32_t block_size,
@@ -279,19 +215,21 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
}
}();
kv_cache_t* __restrict__ b_tile_2 = b_tile;
kv_cache_t* __restrict__ b_tile_3 = [&]() {
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
if constexpr (phase == AttentionGemmPhase::QK) {
// k_cache is prepacked
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
} else if constexpr (phase == AttentionGemmPhase::PV) {
// v_cache is prepacked
return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
} else {
TORCH_CHECK(false, "Unreachable");
}
}();
kv_cache_t* __restrict__ b_tile_4 =
c10::BFloat16* __restrict__ b_tile_4 =
b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
kv_cache_t* __restrict__ b_tile_5 =
c10::BFloat16* __restrict__ b_tile_5 =
b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
int64_t b_stride = AMX_TILE_ROW_BYTES;
@@ -312,25 +250,16 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
_tile_zero(7);
}
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
alignas(64) c10::BFloat16 scratch_4[scratch_elems];
alignas(64) c10::BFloat16 scratch_5[scratch_elems];
for (int32_t k = 0; k < k_group_times; ++k) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
_tile_stream_loadd(2, b_tile_2, b_stride);
_tile_dpbf16ps(6, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
_tile_stream_loadd(3, b_tile_3, b_stride);
_tile_dpbf16ps(7, 0, 3);
_tile_loadd(1, a_tile_1, a_tile_stride);
_tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
_tile_stream_loadd(4, b_tile_4, b_stride);
_tile_dpbf16ps(6, 1, 4);
_tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
_tile_stream_loadd(5, b_tile_5, b_stride);
_tile_dpbf16ps(7, 1, 5);
// update ptrs
@@ -350,13 +279,10 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
}
if (has_tail) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
_tile_stream_loadd(2, b_tile_2, b_stride);
_tile_dpbf16ps(6, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
_tile_stream_loadd(3, b_tile_3, b_stride);
_tile_dpbf16ps(7, 0, 3);
}
@@ -376,25 +302,21 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
_tile_loadconfig(&config);
}
};
} // namespace
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = scalar_t;
using kv_cache_t = kv_cache_scalar_t;
using kv_cache_t = scalar_t;
using logits_buffer_t = float;
using partial_output_buffer_t = float;
using prob_buffer_t = scalar_t;
constexpr static int64_t BlockSizeAlignment =
32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
AMX_TILE_ROW_BYTES /
sizeof(kv_cache_t); // KV token num unit of QK and PV phases
constexpr static int64_t HeadDimAlignment =
2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
constexpr static int64_t MaxQHeadNumPerIteration = 32;
@@ -402,9 +324,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr static ISA ISAType = ISA::AMX;
constexpr static bool scale_on_logits = true;
float k_scale = 1.0f;
float v_scale = 1.0f;
public:
AttentionImpl() : current_q_head_num_(0) {
// Use all columns in AMX tiles
@@ -413,50 +332,21 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
~AttentionImpl() { _tile_release(); }
void init_from_input(const AttentionInput* input) {
if constexpr (fp8_kv) {
k_scale = input->k_scale_fp8;
v_scale = input->v_scale_fp8;
}
}
float get_output_v_scale() const noexcept {
if constexpr (fp8_kv) {
// AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
// Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
// 2^112.
constexpr float bias =
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
return v_scale * bias;
}
return 1.0f;
}
template <template <typename tile_gemm_t> typename attention>
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
if constexpr (fp8_kv) {
// Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
// shifts the exponent bias from FP8 to BF16 (127), so we multiply by
// 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
const float bias =
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
scale *= k_scale * bias;
}
if (q_head_num > AMX_TILE_ROW_NUM) {
if (q_head_num != current_q_head_num_) {
current_q_head_num_ = q_head_num;
TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
amx_tile_config_);
TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
}
attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
attention<TileGemm224<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
} else {
if (q_head_num != current_q_head_num_) {
current_q_head_num_ = q_head_num;
TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
amx_tile_config_);
TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
}
attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
attention<TileGemm122<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
}
}
@@ -521,26 +411,13 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
// reshape KV to AMX friendly layout
static void reshape_and_cache(
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float k_inv = 0.0f, const float v_inv = 0.0f) {
if constexpr (fp8_kv) {
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
key, value, reinterpret_cast<uint8_t*>(key_cache),
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
head_num, head_dim, block_size, key_token_num_stride,
key_head_num_stride, value_token_num_stride, value_head_num_stride,
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
cache_head_num_stride, k_inv, v_inv);
return;
}
const int64_t block_size, const int64_t block_size_stride) {
// For AMX 2D tiles, size of each line is 64 bytes
constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
// For AMX B matrix, N always is 16
@@ -549,9 +426,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
// For now suppose block_size is divisible by amx_tile_column_num
TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -579,7 +453,8 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr int64_t quadword_num_per_group =
token_num_per_group * quadword_num;
int32_t* key_cache_start_ptr =
reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
reinterpret_cast<int32_t*>(key_cache +
block_idx * num_blocks_stride +
head_idx * cache_head_num_stride) +
group_idx * quadword_num_per_group + group_offset;
@@ -608,7 +483,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * value_token_num_stride +
head_idx * value_head_num_stride;
scalar_t* value_cache_start_ptr =
vc + block_idx * num_blocks_stride +
value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride +
sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
sub_group_offset;
-214
View File
@@ -1,214 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#pragma once
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <type_traits>
#include "cpu/utils.hpp"
typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
typedef float __attribute__((__may_alias__)) f32_alias_t;
// Reference scalar dequant — used to verify vectorized AMX dequant.
inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
// NaN encoding in E4M3
if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
uint32_t b_u32 = static_cast<uint32_t>(b);
uint32_t sign = (b_u32 & 0x80) << 24;
uint32_t payload = (b_u32 & 0x7F) << 20;
uint32_t bits = sign | payload;
float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
return b_f32_scaled;
}
inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
v *= inv_scale;
constexpr float fp8_max = 448.0f;
v = std::max(-fp8_max, std::min(fp8_max, v));
if (v == 0.0f) return 0;
// Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
// bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
float v_f32_unscaled = v * 0x1p-120f;
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
if (payload == 0) return sign;
payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding
return static_cast<uint8_t>(sign | payload);
}
// ---------------------------------------------------------------------------
// AMX reshape impl — parameterised on the quantisation function.
// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
// block_size must be divisible by 32.
// ---------------------------------------------------------------------------
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
inline void reshape_and_cache_fp8_amx_impl(
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
float v_inv) {
constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM
const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16
const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
constexpr int64_t head_elems_per_group = 16;
constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16)
const int64_t group_num = head_dim / head_elems_per_group;
const int64_t group_size = block_size * head_elems_per_group;
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t tok = 0; tok < token_num; ++tok) {
for (int64_t h = 0; h < head_num; ++h) {
const int64_t slot = slot_ptr[tok];
if (slot < 0) continue;
const int64_t block_idx = slot / block_size;
const int64_t block_offset = slot % block_size;
// Key: halfword-packed, 2 FP8 per uint16
{
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
const int64_t group_idx = block_offset / token_num_per_group;
const int64_t group_offset = block_offset % token_num_per_group;
uint16_t* kdst =
reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
h * kc_stride1) +
group_idx * halfword_num_per_group + group_offset;
for (int64_t j = 0; j < halfword_num; ++j) {
uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
uint8_t bytes[2] = {fp8_0, fp8_1};
uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
kdst[j * token_num_per_group] = hw;
}
}
// Value: sub-group packing (token_num_per_sub_group = 2)
{
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
uint8_t* vdst =
value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
sub_group_idx * token_num_per_sub_group * head_elems_per_group +
sub_group_offset;
for (int64_t i = 0; i < group_num; ++i) {
for (int64_t j = 0; j < head_elems_per_group; ++j)
vdst[j * token_num_per_sub_group] =
quant_fn(static_cast<float>(vsrc[j]), v_inv);
vsrc += head_elems_per_group;
vdst += group_size;
}
}
}
}
}
// ---------------------------------------------------------------------------
// FP8 E5M2 scalar helpers
// ---------------------------------------------------------------------------
// Reference scalar dequant — used to verify vectorized AMX dequant.
// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
// Byte b → FP16 bits = b << 8 (no bias correction needed).
inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
const uint8_t exp_bits = (b >> 2) & 0x1F;
const uint8_t mant_bits = b & 0x03;
// NaN: exp=11111, mant!=00
if (exp_bits == 0x1F && mant_bits != 0)
return std::numeric_limits<float>::quiet_NaN();
const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
if (exp_bits == 0x1F)
return sign ? -std::numeric_limits<float>::infinity()
: std::numeric_limits<float>::infinity();
if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
if (mant_bits == 0) return 0.0f;
float v = mant_bits * 0x1p-16f;
return (sign ? -v : v) * scale;
}
// Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
uint32_t fp32_bits = sign |
((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
(static_cast<uint32_t>(mant_bits) << 21);
float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
return val * scale;
}
inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
v *= inv_scale;
constexpr float fp8_e5m2_max = 57344.0f;
v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
if (v == 0.0f) return 0;
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
if (exp_fp32 < -14) { // subnormal in E5M2
const int shift = -14 - exp_fp32;
if (shift + 21 >= 32)
return sign; // underflow: too small for E5M2 subnormal
const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
}
const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
return sign | (exp5 << 2) | mant2;
}
// ---------------------------------------------------------------------------
// Select the FP8 quant function at compile time based on kv_cache_t.
// ---------------------------------------------------------------------------
template <typename kv_cache_t>
constexpr auto select_fp8_quant_fn() {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
return float_to_fp8e5m2_scalar;
else
return float_to_fp8e4m3_scalar;
}
// ---------------------------------------------------------------------------
// VEC reshape impl — parameterised on the quantisation function.
// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
// #pragma cannot appear inside variadic macro arguments.
// ---------------------------------------------------------------------------
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
inline void reshape_and_cache_fp8_vec_impl(
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
float v_inv) {
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t tok = 0; tok < token_num; ++tok) {
for (int64_t h = 0; h < head_num; ++h) {
const int64_t slot = slot_ptr[tok];
if (slot < 0) continue;
const int64_t block_idx = slot / block_size;
const int64_t block_offset = slot % block_size;
// Key layout: column-major within block
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
block_offset;
for (int64_t i = 0; i < head_dim; ++i)
kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
// Value layout: row-major within block (contiguous head_dim bytes)
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
h * vc_stride1 + block_offset * head_dim;
for (int64_t i = 0; i < head_dim; ++i)
vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
}
}
}
+7 -31
View File
@@ -14,22 +14,8 @@
namespace cpu_attention {
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
enum class Fp8KVCacheDataType {
kAuto = 0,
kFp8E4M3 = 1,
kFp8E5M2 = 2,
};
struct AttentionInput;
template <ISA isa, typename scalar_t, int64_t head_dim,
typename kv_cache_scalar_t = scalar_t>
class AttentionImpl {
public:
void init_from_input(const AttentionInput*) {}
float get_output_v_scale() const noexcept { return 1.0f; }
};
template <ISA isa, typename scalar_t, int64_t head_dim>
class AttentionImpl {};
struct AttentionWorkItemGroup {
int32_t req_id;
@@ -794,9 +780,6 @@ struct AttentionInput {
int32_t sliding_window_left;
int32_t sliding_window_right;
float softcap;
// FP8 KV cache scales (used by FP8 attention implementations)
float k_scale_fp8 = 1.0f;
float v_scale_fp8 = 1.0f;
};
#define DEFINE_CPU_ATTENTION_PARAMS \
@@ -1391,13 +1374,6 @@ class AttentionMainLoop {
}
attention_impl_t attn_impl;
constexpr bool fp8_kv = std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
float output_v_scale = 1.0f;
if constexpr (fp8_kv) {
attn_impl.init_from_input(input);
output_v_scale = attn_impl.get_output_v_scale();
}
// general information
const int32_t q_head_num = input->num_heads;
@@ -1777,7 +1753,7 @@ class AttentionMainLoop {
reinterpret_cast<query_t*>(input->output) +
output_buffer_offset,
sum_buffer, actual_q_heads_per_kv,
actual_q_token_num, q_head_num, output_v_scale);
actual_q_token_num, q_head_num);
} else {
const int32_t stride =
actual_q_heads_per_kv * split_kv_q_token_num_threshold;
@@ -1847,7 +1823,7 @@ class AttentionMainLoop {
split_output_buffer,
reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
q_head_num, output_v_scale);
q_head_num);
}
}
}
@@ -1971,8 +1947,8 @@ class AttentionMainLoop {
query_t* __restrict__ curr_output_buffer,
float* __restrict__ sum_buffer,
const int32_t q_heads_per_kv,
const int32_t actual_q_token_num, const int32_t q_head_num,
const float v_scale = 1.0f) {
const int32_t actual_q_token_num,
const int32_t q_head_num) {
// final output
using output_vec_t = typename VecTypeTrait<query_t>::vec_t;
@@ -1986,7 +1962,7 @@ class AttentionMainLoop {
curr_partial_output_buffer;
query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
vec_op::FP32Vec16 inv_sum_scale_vec(v_scale / *curr_sum_buffer);
vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
for (int32_t i = 0; i < group_num_per_head; ++i) {
vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
+4 -5
View File
@@ -248,8 +248,8 @@ class TileGemmNeonFMLA {
} // namespace
// this is similar to "ISA::VEC" at the moment
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
@@ -343,8 +343,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -389,7 +388,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
#ifdef ARM_BF16_SUPPORT
// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
template <int64_t head_dim>
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
head_dim> {};
#endif
+1 -2
View File
@@ -602,8 +602,7 @@ class AttentionImplNEONBFMMLA {
[[maybe_unused]] const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size,
[[maybe_unused]] const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
[[maybe_unused]] const int64_t block_size_stride) {
const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
const int64_t v_pair_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+28 -105
View File
@@ -1,37 +1,11 @@
#ifndef CPU_ATTN_VEC_HPP
#define CPU_ATTN_VEC_HPP
#include "cpu_attn_fp8.hpp"
#include "cpu_attn_impl.hpp"
namespace cpu_attention {
namespace {
// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
// covering the lower 16 and upper 16 positions.
// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
template <typename kv_cache_t>
FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
const kv_cache_t* ptr) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
// BF16 container, but values are in the FP16 exponent range (bias 15 not
// 127).
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
vec_op::fp8_e4m3_tag{});
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
} else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
vec_op::fp8_e5m2_tag{});
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
} else {
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
return {vec_op::FP32Vec16(load_vec_t(ptr)),
vec_op::FP32Vec16(load_vec_t(ptr + 16))};
}
}
// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
template <typename kv_cache_t>
class TileGemm82 {
@@ -80,7 +54,10 @@ class TileGemm82 {
const int32_t block_size, const int32_t dynamic_k_size,
const bool accum_c) {
static_assert(0 < M && M <= 8);
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
kv_cache_t* __restrict__ curr_b_0 = b_tile;
kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
float* __restrict__ curr_c_0 = c_tile;
float* __restrict__ curr_c_1 = c_tile + 16;
@@ -99,14 +76,16 @@ class TileGemm82 {
}
float* __restrict__ curr_a = a_tile;
kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
load_vec_t b_0_reg(curr_b_0);
vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
load_vec_t b_1_reg(curr_b_1);
vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
vec_op::FP32Vec16 a_reg(*curr_m_a);
float v = *curr_m_a;
vec_op::FP32Vec16 a_reg(v);
c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
@@ -116,7 +95,8 @@ class TileGemm82 {
// update
curr_a += 1;
curr_b += ldb;
curr_b_0 += ldb;
curr_b_1 += ldb;
}
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -129,20 +109,15 @@ class TileGemm82 {
});
}
};
} // namespace
// This is a general but naive implementation based on vector instructions
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
using kv_cache_t = kv_cache_scalar_t;
using kv_cache_t = scalar_t;
using logits_buffer_t = float;
using partial_output_buffer_t = float;
using prob_buffer_t = float;
@@ -154,45 +129,11 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr static int64_t MaxQHeadNumPerIteration = 8;
constexpr static int64_t HeadDim = head_dim;
constexpr static ISA ISAType = ISA::VEC;
constexpr static bool scale_on_logits = fp8_kv;
float k_scale = 1.0f;
float v_scale = 1.0f;
constexpr static bool scale_on_logits = false; // apply scale on q_buffer
public:
void init_from_input(const AttentionInput* input) {
if constexpr (fp8_kv) {
k_scale = input->k_scale_fp8;
v_scale = input->v_scale_fp8;
}
}
float get_output_v_scale() const noexcept {
if constexpr (fp8_kv) {
// VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
// E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
// so no correction.
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
return v_scale;
} else {
return v_scale * 0x1p8f;
}
}
return 1.0f;
}
template <template <typename tile_gemm_t> typename attention>
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
if constexpr (fp8_kv) {
// Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
// uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
// correction.
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
scale *= k_scale;
} else {
scale *= k_scale * 0x1p8f;
}
}
attention<TileGemm82<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
}
@@ -220,19 +161,17 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
// row-major
}
// Copy q to q_buffer and cast it to fp32.
// FP8: QK scale is folded into execute_attention; copy Q unscaled here.
void copy_q_heads_tile(scalar_t* __restrict__ src,
float* __restrict__ q_buffer, const int32_t q_num,
const int32_t q_heads_per_kv,
const int64_t q_num_stride,
const int64_t q_head_stride, float scale) {
// Copy q to q_buffer and cast it to fp32
static void copy_q_heads_tile(
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
float* __restrict__ q_buffer, const int32_t q_num,
const int32_t q_heads_per_kv, const int64_t q_num_stride,
const int64_t q_head_stride, float scale) {
static_assert(head_dim % 16 == 0);
constexpr int32_t unroll_size = head_dim / 16;
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
const float effective_scale = fp8_kv ? 1.0f : scale;
vec_op::FP32Vec16 scale_vec(effective_scale);
vec_op::FP32Vec16 scale_vec(scale);
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
scalar_t* __restrict__ curr_q =
@@ -257,26 +196,13 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
// reshape K as column-major and V as row-major
static void reshape_and_cache(
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float k_inv = 0.0f, const float v_inv = 0.0f) {
if constexpr (fp8_kv) {
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
key, value, reinterpret_cast<uint8_t*>(key_cache),
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
head_num, head_dim, block_size, key_token_num_stride,
key_head_num_stride, value_token_num_stride, value_head_num_stride,
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
cache_head_num_stride, k_inv, v_inv);
return;
}
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -294,9 +220,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * key_token_num_stride +
head_idx * key_head_num_stride;
scalar_t* key_cache_start_ptr =
reinterpret_cast<scalar_t*>(key_cache) +
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
block_offset;
key_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride + block_offset;
#pragma GCC unroll 8
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
@@ -309,9 +234,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * value_token_num_stride +
head_idx * value_head_num_stride;
scalar_t* value_cache_start_ptr =
reinterpret_cast<scalar_t*>(value_cache) +
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
block_offset * head_dim;
value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride + block_offset * head_dim;
std::memcpy(value_cache_start_ptr, value_start_ptr,
sizeof(scalar_t) * head_dim);
}
@@ -319,7 +243,6 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
}
}
};
} // namespace cpu_attention
#endif
+3 -3
View File
@@ -116,9 +116,9 @@ class TileGemm161 {
} // namespace
// This is a general but naive implementation based on vector instructions
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
: public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
: public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
+3 -4
View File
@@ -244,8 +244,8 @@ class TileGemmS390X {
} // namespace
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
@@ -342,8 +342,7 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-6
View File
@@ -15,9 +15,6 @@ using namespace at::vec;
namespace vec_op {
struct fp8_e4m3_tag {};
struct fp8_e5m2_tag {};
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -325,9 +322,6 @@ struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
reg.val[2] = vec8_data.reg.val[0];
reg.val[3] = vec8_data.reg.val[0];
};
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
};
struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
-6
View File
@@ -8,9 +8,6 @@
#include <torch/all.h>
namespace vec_op {
struct fp8_e4m3_tag {};
struct fp8_e5m2_tag {};
#define vec_neg(a) (-(a))
#define vec_add(a, b) ((a) + (b))
#define vec_sub(a, b) ((a) - (b))
@@ -244,9 +241,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
explicit BF16Vec32(const BF16Vec8& vec8_data)
: reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
};
-139
View File
@@ -11,17 +11,6 @@ static_assert(false, "AVX2 must be supported for the current implementation.");
namespace vec_op {
// Tags for FP8 BF16Vec32 constructors (avoid overload collision with
// BF16Vec32(void*)).
// VEC path (FP8 → pseudo-FP16 layout, scale correction applied later):
struct fp8_e4m3_tag {}; // E4M3 → pseudo-FP16; BF16 value = true_E4M3 * 2^-8
struct fp8_e5m2_tag {}; // E5M2 → FP16 bits directly (same exponent bias=15)
// AMX path (FP8 → unscaled BF16, no FP32 round-trip):
// BF16 value = true_E4M3 * 2^-120 (E4M3) or true_E5M2 * 2^-112 (E5M2).
// Exponent rebiasing is folded into k/v scales by the caller.
struct fp8_bf16_e4m3_tag {};
struct fp8_bf16_e5m2_tag {};
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
@@ -187,50 +176,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
(__m128i)vec8_data.reg, 2),
(__m128i)vec8_data.reg, 3)) {}
// Decode 32 FP8-E4M3 bytes to pseudo-FP16 layout (stored in the BF16
// register). Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 7);
reg = _mm512_or_si512(sign, payload);
}
// Decode 32 FP8-E5M2 bytes to FP16 layout.
// E5M2 and FP16 share the same 5-bit exponent bias (15), so FP8 byte b maps
// directly to FP16 bits by shifting left 8 — no sign/payload reconstruction.
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
reg = _mm512_slli_epi16(_mm512_cvtepu8_epi16(b8), 8);
}
// Direct FP8-E4M3 → unscaled BF16 for AMX (no FP32 round-trip).
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 4);
reg = _mm512_or_si512(sign, payload);
}
// Direct FP8-E5M2 → unscaled BF16 for AMX (no FP32 round-trip).
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 5);
reg = _mm512_or_si512(sign, payload);
}
void save(void* ptr) const { *reinterpret_cast<__m512i*>(ptr) = reg; }
};
#else
@@ -255,77 +200,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
_mm256_castsi128_si256((__m128i)vec8_data.reg),
(__m128i)vec8_data.reg, 1)) {}
// E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
// above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
__m256i sign_low = _mm256_slli_epi16(
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)), 8);
__m256i payload_low = _mm256_slli_epi16(
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)), 7);
__m256i sign_high = _mm256_slli_epi16(
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)), 8);
__m256i payload_high = _mm256_slli_epi16(
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)), 7);
reg_low = _mm256_or_si256(sign_low, payload_low);
reg_high = _mm256_or_si256(sign_high, payload_high);
}
// E5M2 decode (AVX2 path) — b << 8 maps to FP16 bits; see AVX512 variant
// above.
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
reg_low = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_low), 8);
reg_high = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_high), 8);
}
// Direct FP8-E4M3 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
reg_low = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
4));
reg_high = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
4));
}
// Direct FP8-E5M2 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
reg_low = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
5));
reg_high = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
5));
}
void save(void* ptr) const {
_mm256_storeu_si256((__m256i*)ptr, reg_low);
_mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
@@ -516,11 +390,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
: reg(_mm512_castsi512_ps(
_mm512_bslli_epi128(_mm512_cvtepu16_epi32(v.reg), 2))) {}
explicit FP32Vec16(const BF16Vec32& v, int upper) {
__m256i v_half_i = _mm512_extracti32x8_epi32(v.reg, upper);
reg = _mm512_cvtph_ps(v_half_i);
}
explicit FP32Vec16(const FP16Vec16& v) : reg(_mm512_cvtph_ps(v.reg)) {}
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -625,14 +494,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(const FP32Vec8& data)
: reg_low(data.reg), reg_high(data.reg) {}
explicit FP32Vec16(const BF16Vec32& v, int upper) {
const __m256i& half = upper ? v.reg_high : v.reg_low;
__m128i lo = _mm256_extractf128_si256(half, 0);
__m128i hi = _mm256_extractf128_si256(half, 1);
reg_low = _mm256_cvtph_ps(lo);
reg_high = _mm256_cvtph_ps(hi);
}
explicit FP32Vec16(const FP16Vec16& v) {
__m128i low = _mm256_extractf128_si256(v.reg, 0);
__m128i high = _mm256_extractf128_si256(v.reg, 1);
+125 -141
View File
@@ -22,95 +22,71 @@ ISA_TYPES = {
"VXE": 4,
}
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
KV_CACHE_IDX = {
"auto": 0,
"fp8_e4m3": 1,
"fp8_e5m2": 2,
}
# C++ type for each kv_cache index
KV_CACHE_CPP_TYPES = {
"auto": "scalar_t",
"fp8_e4m3": "c10::Float8_e4m3fn",
"fp8_e5m2": "c10::Float8_e5m2",
}
# ISAs supported for head_dims divisible by 32
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
# ISAs supported for head_dims divisible by 16 only
ISA_FOR_16 = ["VEC16"]
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
ISA_FOR_FP8 = ["AMX", "VEC"]
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
"""Encode head_dim, ISA type, and KV cache type into a single int64_t."""
def encode_params(head_dim: int, isa_type: str) -> int:
"""Encode head_dim and ISA type into a single int64_t."""
isa_val = ISA_TYPES[isa_type]
kv_val = KV_CACHE_IDX[kv_cache]
# Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
# This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
return (head_dim << 16) | (kv_val << 8) | isa_val
# Encoding: (head_dim << 8) | isa_type
# This allows head_dim up to 2^56 - 1 and 256 ISA types
return (head_dim << 8) | isa_val
def _make_case(
head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
) -> str:
"""Generate a single switch case line."""
encoded = encode_params(head_dim, isa, kv_cache)
actual_isa = isa_override if isa_override else isa
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
attn_impl = (
f"cpu_attention::AttentionImpl<"
f"cpu_attention::ISA::{actual_isa}, \\\n"
f" "
f"scalar_t, head_dim, {cpp_type}>"
)
comment = (
f"head_dim={head_dim}, isa={isa}"
if kv_cache == "auto"
else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
)
return (
f""" case {encoded}LL: {{ """
f"""/* {comment} */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = {attn_impl}; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
def generate_cases_for_isa_group(isa_list: list[str]) -> str:
"""Generate switch cases for a specific ISA group."""
cases = []
# Non-FP8 cases for head_dims divisible by 32
# Generate cases for head_dims divisible by 32
for head_dim in HEAD_DIMS_32:
for isa in isa_list:
if isa not in ISA_FOR_32:
continue
cases.append(_make_case(head_dim, isa, "auto"))
encoded = encode_params(head_dim, isa)
case_str = (
f""" case {encoded}LL: {{ """
f"""/* head_dim={head_dim}, isa={isa} */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = cpu_attention::AttentionImpl<"""
f"""cpu_attention::ISA::{isa}, \\"""
f"""
"""
f"""scalar_t, head_dim>; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
cases.append(case_str)
# Non-FP8 cases for head_dims divisible by 16 only
# Generate cases for head_dims divisible by 16 only
for head_dim in HEAD_DIMS_16:
for isa in isa_list:
cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
# FP8 cases: only AMX and VEC, only head_dims divisible by 32
if include_fp8:
for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
for head_dim in HEAD_DIMS_32:
for isa in isa_list:
if isa not in ISA_FOR_FP8:
continue
cases.append(_make_case(head_dim, isa, fp8_type))
encoded = encode_params(head_dim, isa)
case_str = (
f""" case {encoded}LL: {{ """
f"""/* head_dim={head_dim}, isa={isa} """
f"""(using VEC16) */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = cpu_attention::AttentionImpl<"""
f"""cpu_attention::ISA::VEC16, \\"""
f"""
"""
f"""scalar_t, head_dim>; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
cases.append(case_str)
return "\n".join(cases)
@@ -118,9 +94,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)
def generate_helper_function() -> str:
"""Generate helper function to encode parameters."""
return """
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
int64_t kv_cache_idx = 0) {
return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
return (head_dim << 8) | static_cast<int64_t>(isa);
}
"""
@@ -154,78 +129,87 @@ def generate_header_file() -> str:
# Generate dispatch macro with conditional compilation for different ISA sets
header += """
// Dispatch macro using encoded parameters.
// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
// Dispatch macro using encoded parameters
"""
def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
"""Return one CPU_ATTN_DISPATCH macro block for a given guard."""
enc = (
" int64_t encoded_params = encode_cpu_attn_params("
"HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
)
cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
tail = (
"\n"
" default: { \\\n"
" TORCH_CHECK(false, "
'"Unsupported CPU attention configuration: head_dim=" + \\\n'
' std::to_string(HEAD_DIM) + " isa=" + \\\n'
" std::to_string(static_cast<int>(ISA_TYPE))"
" + \\\n"
' " kv_cache_idx=" + '
"std::to_string(KV_CACHE_IDX)); \\\n"
" } \\\n"
" } \\\n"
" }()\n\n"
)
return (
f"{guard}\n"
"#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
" [&] { \\\n"
f"{enc}\n"
" switch (encoded_params) { \\\n"
f"{cases}"
f"{tail}"
)
# x86_64 with AMX
header += """#if defined(CPU_CAPABILITY_AMXBF16)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
header += _macro_block(
"#if defined(CPU_CAPABILITY_AMXBF16)",
["AMX", "VEC", "VEC16"],
fp8=True,
)
header += _macro_block(
"#elif defined(__aarch64__)",
["NEON", "VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#elif defined(__s390x__)",
["VXE", "VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#elif defined(__AVX512F__)",
["VEC", "VEC16"],
fp8=True,
)
header += _macro_block(
"#elif defined(__AVX2__)",
["VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#else",
["VEC", "VEC16"],
fp8=False,
)
header += (
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
)
"""
# ARM64 with NEON
header += """#elif defined(__aarch64__)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
"""
# s390x with VXE
header += """#elif defined(__s390x__)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
"""
# Fallback: VEC and VEC16 only
header += """#else
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
#endif // CPU_ATTN_DISPATCH_GENERATED_H
"""
return header
+5 -11
View File
@@ -101,9 +101,7 @@ void cpu_attn_reshape_and_cache(const torch::Tensor& key,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
const torch::Tensor& slot_mapping,
const std::string& isa, const double k_scale,
const double v_scale,
const std::string& kv_cache_dtype);
const std::string& isa);
void cpu_attention_with_kv_cache(
const torch::Tensor& query, const torch::Tensor& key_cache,
@@ -114,8 +112,7 @@ void cpu_attention_with_kv_cache(
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, const double softcap,
const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, const double k_scale,
const double v_scale, const std::string& kv_cache_dtype);
const std::optional<torch::Tensor>& s_aux);
// Note: just for avoiding importing errors
void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
@@ -387,18 +384,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
&get_scheduler_metadata);
ops.def(
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
"()",
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
"isa) -> ()",
&cpu_attn_reshape_and_cache);
ops.def(
"cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
"value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
"seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
"sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
"()",
"float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
&cpu_attention_with_kv_cache);
// placeholders
+35 -10
View File
@@ -96,14 +96,44 @@ struct enable_sm90_or_later : Kernel {
};
template <typename Kernel>
struct enable_sm100_to_sm120 : Kernel {
struct enable_sm90_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
#if __CUDA_ARCH__ == 900
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm[100, 120).\n");
printf("This kernel only supports sm90.\n");
asm("trap;");
#endif
#endif
}
};
template <typename Kernel>
struct enable_sm100f_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm100f.\n");
asm("trap;");
#endif
#endif
}
};
template <typename Kernel>
struct enable_sm100a_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if __CUDA_ARCH__ == 1000
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm100a.\n");
asm("trap;");
#endif
#endif
@@ -118,7 +148,7 @@ struct enable_sm120_only : Kernel {
#if __CUDA_ARCH__ == 1200
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm120a.\n");
printf("This kernel only supports sm120.\n");
asm("trap;");
#endif
#endif
@@ -130,13 +160,8 @@ template <typename Kernel>
struct enable_sm120_family : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm120f.\n");
asm("trap;");
#endif
#endif
}
};
+693
View File
@@ -0,0 +1,693 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// DeepSeek V4 indexer top-k (k = 512 for Flash, k = 1024 for Pro). Ported
// from sglang's jit_kernel/csrc/deepseek_v4/topk_v2.cuh.
//
// Combines three strategies (Register / Streaming / Cluster) dispatched per
// row by a separate plan kernel that decides a `cluster_threshold` from the
// observed seq_lens distribution. The host side picks one of three launch
// shapes:
// 1. all rows fit in the small (register) path -> single short kernel
// 2. small batch (<= kNumClusters) with some long rows -> fused cluster
// kernel (stage 1 + tie-break in one launch)
// 3. larger batch -> persistent cluster stage 1 + non-cluster stage 2
//
// Architecture support: Hopper (sm_90a) and Blackwell datacenter (sm_100/
// sm_103). Requires thread-block clusters, TMA bulk async copy, mbarrier,
// and Programmatic Dependent Launch — sm_120 (consumer Blackwell) lacks
// clusters and is not supported. The heuristic constants in `topk_plan`
// were tuned on B200 (sglang upstream); they are functionally correct on
// H100/H200 too but may be suboptimal until retuned.
#include "topk/cluster.cuh"
#include "topk/common.cuh"
#include "topk/register.cuh"
#include "topk/streaming.cuh"
#include "topk/utils.cuh"
#include "core/registration.h"
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/Exception.h>
#include <cooperative_groups.h>
#include <cuda_runtime.h>
#include <torch/all.h>
#include <torch/library.h>
#include <algorithm>
#include <cstdint>
namespace vllm::dsv4_topk {
// All K-dependent type and constant lookups go through these aliases / vars
// so the kernels can be templated on K. Kernel and Smem sizes happen to be
// K-independent (e.g., kMaxTies, kMax2PassLength, kHistBins are all set in
// terms of kBlockSize/kHistBits, not K), so we don't pay extra smem for the
// 1024 instantiation.
template <uint32_t K> using Large = ClusterTopK<K>;
template <uint32_t K> using Medium = StreamingTopK<K>;
template <uint32_t K> using Small = RegisterTopK<K>;
// Metadata struct layout is K-independent — pick any K to grab the type.
using Metadata = Large<512>::Metadata;
constexpr uint32_t kNumClusters = 15; // hardware-capped persistent count
constexpr uint32_t kClusterSize = Large<512>::kClusterSize;
constexpr uint32_t kMax2PassLength = Small<512>::kMax2PassLength;
constexpr uint32_t kMaxSupportedLength = Large<512>::kMaxLength;
// Row 0 of the metadata tensor stores GlobalMetadata; rows [1..N+1) hold the
// per-item Metadata entries that the persistent stage-1 consumes.
struct alignas(16) GlobalMetadata {
uint32_t cluster_threshold;
uint32_t num_cluster_items;
uint32_t reserved[2];
};
static_assert(sizeof(GlobalMetadata) == sizeof(Metadata),
"metadata row 0 layout must match Metadata stride");
#define VLLM_SMALL_TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
#define VLLM_LARGE_CLUSTER __cluster_dims__(1, kClusterSize, 1)
// Stage 1 is persistent + cluster -> high smem -> occupancy 1.
#define VLLM_LARGE_TOPK_STAGE_1 \
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
// Stage 2 is non-cluster + small smem -> occupancy 2.
#define VLLM_LARGE_TOPK_STAGE_2 __global__ __launch_bounds__(kBlockSize, 2)
#define VLLM_FUSED_COMBINE_KERNEL \
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
#define VLLM_PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
struct TopKParams {
const uint32_t* __restrict__ seq_lens;
const float* __restrict__ scores;
const int32_t* __restrict__ page_table;
int32_t* __restrict__ page_indices;
int64_t score_stride;
int64_t page_table_stride;
uint8_t* __restrict__ workspace;
const Metadata* __restrict__ metadata = nullptr;
int64_t workspace_stride; // bytes per batch
uint32_t batch_size;
uint32_t page_bits;
VLLM_DSV4_DEVICE const float* get_scores(uint32_t batch_id) const {
return scores + batch_id * score_stride;
}
template <uint32_t K, bool kRawOutput>
VLLM_DSV4_DEVICE TransformParamsT<kRawOutput> get_transform(
uint32_t batch_id, int32_t* indices) const {
return {
.page_table = page_table + batch_id * page_table_stride,
.indices_in = indices,
.indices_out = page_indices + batch_id * K,
.page_bits = page_bits,
};
}
VLLM_DSV4_DEVICE const GlobalMetadata& get_global_metadata() const {
return *reinterpret_cast<const GlobalMetadata*>(metadata);
}
VLLM_DSV4_DEVICE const Metadata& get_item_metadata(uint32_t work_id) const {
return metadata[1 + work_id]; // skip the GlobalMetadata row
}
};
VLLM_DSV4_DEVICE uint2 partition_work(uint32_t length, uint32_t rank) {
constexpr uint32_t kTMAAlign = 4;
const auto total_units = (length + kTMAAlign - 1) / kTMAAlign;
const auto base = total_units / kClusterSize;
const auto extra = total_units % kClusterSize;
const auto local_units = base + (rank < extra ? 1u : 0u);
const auto offset_units = rank * base + min(rank, extra);
const auto offset = offset_units * kTMAAlign;
const auto finish = min(offset + local_units * kTMAAlign, length);
return {offset, finish - offset};
}
// --------------------------------------------------------------------------
// Plan kernel: decides cluster_threshold from the observed seq_lens
// distribution and compacts items with seq_len > threshold into metadata[1..].
// --------------------------------------------------------------------------
VLLM_PLAN_KERNEL void topk_plan(const uint32_t* __restrict__ seq_lens,
Metadata* __restrict__ metadata,
uint32_t batch_size,
uint32_t static_cluster_threshold) {
// (threshold, max_batch_size_for_that_threshold). Tuned on B200 by sglang.
struct Pair {
uint32_t threshold;
uint32_t max_batch_size;
};
constexpr Pair kCandidates[] = {
{32768, 30}, {40960, 45}, {49152, 45}, {65536, 60},
{98304, 60}, {131072, 75}, {196608, 90}, {262144, 105},
};
constexpr uint32_t kNumCandidates =
sizeof(kCandidates) / sizeof(kCandidates[0]);
constexpr uint32_t kMinBatchSize = kCandidates[0].max_batch_size;
static_assert(kCandidates[0].threshold == kMax2PassLength);
static_assert(kCandidates[kNumCandidates - 1].threshold ==
kMaxSupportedLength);
__shared__ uint32_t s_count;
__shared__ uint32_t s_counts[kNumCandidates];
__shared__ uint32_t s_threshold;
const auto tx = threadIdx.x;
if (tx == 0) s_count = 0;
if (tx < kNumCandidates) s_counts[tx] = 0;
__syncthreads();
if (static_cluster_threshold > 0) {
if (tx == 0) s_threshold = static_cluster_threshold;
} else if (batch_size <= kMinBatchSize) {
if (tx == 0) s_threshold = kMax2PassLength;
} else {
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
assert(sl <= kMaxSupportedLength);
uint32_t count = 0;
#pragma unroll
for (uint32_t j = 0; j < kNumCandidates; ++j) {
count += (sl > kCandidates[j].threshold ? 1 : 0);
}
if (count > 0) {
atomicAdd(&s_counts[count - 1], 1);
}
}
__syncthreads();
if (tx == 0) {
uint32_t accum = 0;
uint32_t chosen = kMaxSupportedLength;
#pragma unroll
for (uint32_t i = 0; i < kNumCandidates; ++i) {
const auto j = kNumCandidates - 1 - i;
accum += s_counts[j];
if (accum > kCandidates[j].max_batch_size) break;
chosen = kCandidates[j].threshold;
}
s_threshold = chosen;
}
}
__syncthreads();
const auto cluster_threshold = max(s_threshold, kMax2PassLength);
// Compact items with seq_len > cluster_threshold into metadata[1..N+1).
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
const uint32_t sl = seq_lens[i];
if (sl > cluster_threshold) {
const auto pos = atomicAdd(&s_count, 1);
metadata[1 + pos] = {i, sl, false};
}
}
__syncthreads();
const auto N = s_count;
// has_next chain for the persistent consumer + sentinel slots.
for (uint32_t i = tx; i < N; i += kBlockSize) {
if (i + kNumClusters < N) metadata[1 + i].has_next = true;
}
if (tx < kNumClusters && tx >= N) metadata[1 + tx] = {0, 0, false};
if (tx == 0) {
auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
*g = {
.cluster_threshold = cluster_threshold,
.num_cluster_items = N,
.reserved = {0, 0},
};
}
}
// --------------------------------------------------------------------------
// Short kernel: all rows fit in the register path (max_seq_len <=
// Small::kMax1PassLength).
// --------------------------------------------------------------------------
template <uint32_t K, bool kRawOutput>
VLLM_SMALL_TOPK_KERNEL void topk_short_transform(
const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto seq_len = params.seq_lens[batch_id];
const auto transform =
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
if (seq_len <= K) {
trivial_transform(transform, seq_len, K);
} else {
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
/*use_pdl=*/true);
pdl_trigger_secondary<true>();
Small<K>::transform(transform);
}
}
// --------------------------------------------------------------------------
// Persistent stage 1 (cluster). One CTA per cluster; the persistent block
// walks `metadata[1..N]` round-robin and runs Large::stage1 per item.
// --------------------------------------------------------------------------
template <uint32_t K, bool kRawOutput>
VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess(
const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
uint32_t work_id = blockIdx.x;
uint32_t batch_id = 0, seq_len = 0, length = 0, offset = 0;
bool has_next = false;
const auto cluster_rank = blockIdx.y;
const auto prefetch_metadata = [&] {
const auto m = params.get_item_metadata(work_id);
batch_id = m.batch_id;
seq_len = m.seq_len;
has_next = m.has_next;
work_id += kNumClusters;
};
const auto launch_prologue = [&] {
const auto partition = partition_work(seq_len, cluster_rank);
offset = partition.x;
length = partition.y;
Large<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
smem);
};
pdl_wait_primary<true>();
pdl_trigger_secondary<true>();
prefetch_metadata();
if (seq_len == 0) return;
Large<K>::stage1_init(smem);
launch_prologue();
while (true) {
const auto this_length = length;
const auto this_offset = offset;
const auto need_prefetch = has_next;
const auto transform =
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
const auto ws = params.workspace + batch_id * params.workspace_stride;
if (need_prefetch) prefetch_metadata();
Large<K>::stage1(s_topk_indices, this_length, smem, /*reuse=*/true);
if (need_prefetch) launch_prologue();
Large<K>::stage1_epilogue(transform, this_offset, ws, smem);
if (!need_prefetch) break;
}
}
// --------------------------------------------------------------------------
// Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large.
// --------------------------------------------------------------------------
template <uint32_t K, bool kRawOutput>
VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform(
const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto seq_len = params.seq_lens[batch_id];
const auto cluster_threshold = params.get_global_metadata().cluster_threshold;
const auto transform =
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
if (seq_len <= K) {
trivial_transform(transform, seq_len, K);
} else if (seq_len <= kMax2PassLength) {
if (seq_len <= Small<K>::kMax1PassLength) {
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len,
smem);
} else {
__syncwarp();
Small<K>::template run<true>(params.get_scores(batch_id),
s_topk_indices, seq_len, smem);
}
Small<K>::transform(transform);
} else if (seq_len <= cluster_threshold) {
Medium<K>::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem);
Medium<K>::transform(transform, smem);
} else {
const auto ws = params.workspace + batch_id * params.workspace_stride;
pdl_wait_primary<true>();
Large<K>::transform(transform, ws, smem);
}
}
// --------------------------------------------------------------------------
// Fused kernel for small batches. Both stage 1 and the tie-break run inside
// the same launch; cluster rank 0 finishes the row.
// --------------------------------------------------------------------------
template <uint32_t K, bool kRawOutput>
VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform(
const __grid_constant__ TopKParams params) {
alignas(128) extern __shared__ uint8_t smem[];
__shared__ int32_t s_topk_indices[K];
const auto batch_id = blockIdx.x;
const auto cluster_rank = blockIdx.y;
const auto seq_len = params.seq_lens[batch_id];
const auto transform =
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
if (seq_len <= K) {
if (cluster_rank != 0) return;
trivial_transform(transform, seq_len, K);
} else if (seq_len <= Small<K>::kMax1PassLength) {
if (cluster_rank != 0) return;
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
/*use_pdl=*/true);
Small<K>::transform(transform);
} else {
const auto partition = partition_work(seq_len, cluster_rank);
const auto offset = partition.x;
const auto length = partition.y;
const auto ws = params.workspace + batch_id * params.workspace_stride;
Large<K>::stage1_init(smem);
pdl_wait_primary<true>();
Large<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
smem);
Large<K>::stage1(s_topk_indices, length, smem);
Large<K>::stage1_epilogue(transform, offset, ws, smem);
cooperative_groups::this_cluster().sync();
if (cluster_rank != 0) return;
Large<K>::transform(transform, ws, smem);
}
}
template <uint32_t K> constexpr size_t kStage1SMEM = sizeof(typename Large<K>::Smem) + 128;
template <uint32_t K> constexpr size_t kStage2SMEM =
(sizeof(typename Small<K>::Smem) > sizeof(typename Medium<K>::Smem)
? sizeof(typename Small<K>::Smem)
: sizeof(typename Medium<K>::Smem)) +
128;
// Per-(kernel, smem) memoization: each instantiation has its own static. This
// matters because cudaFuncSetAttribute is per-function and we want it to fire
// exactly once per kernel symbol.
template <auto* f, size_t kSmem>
void setup_kernel_smem_once() {
[[maybe_unused]] static const auto result = [] {
return cudaFuncSetAttribute(reinterpret_cast<const void*>(f),
cudaFuncAttributeMaxDynamicSharedMemorySize,
static_cast<int>(kSmem));
}();
TORCH_CHECK(result == cudaSuccess,
"fast_topk_v2: cudaFuncSetAttribute failed: ",
cudaGetErrorString(result));
}
// --------------------------------------------------------------------------
// Host-side launchers
// --------------------------------------------------------------------------
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_DTYPE(x, t) \
TORCH_CHECK(x.scalar_type() == (t), #x " must be ", #t)
#define CHECK_CONTIG(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
} // namespace vllm::dsv4_topk
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
int64_t static_cluster_threshold) {
using namespace vllm::dsv4_topk;
CHECK_CUDA(seq_lens);
CHECK_CUDA(metadata);
CHECK_DTYPE(seq_lens, torch::kInt32);
CHECK_DTYPE(metadata, torch::kInt32);
TORCH_CHECK(seq_lens.dim() == 1);
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4);
TORCH_CHECK(metadata.size(0) == seq_lens.size(0) + 1,
"metadata must be (batch_size + 1, 4)");
CHECK_CONTIG(seq_lens);
CHECK_CONTIG(metadata);
const auto batch_size = static_cast<uint32_t>(seq_lens.size(0));
if (batch_size <= kNumClusters) return; // metadata unused in fused path
const auto stream = at::cuda::getCurrentCUDAStream().stream();
cudaLaunchConfig_t cfg{};
cfg.gridDim = dim3(1);
cfg.blockDim = dim3(kBlockSize);
cfg.dynamicSmemBytes = 0;
cfg.stream = stream;
cfg.numAttrs = 0;
TORCH_CHECK(cudaLaunchKernelEx(
&cfg, &topk_plan,
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
reinterpret_cast<Metadata*>(metadata.data_ptr<int32_t>()),
batch_size,
static_cast<uint32_t>(static_cluster_threshold)) == cudaSuccess,
"fast_topk_v2_plan launch failed: ",
cudaGetErrorString(cudaGetLastError()));
}
namespace vllm::dsv4_topk {
// Shared dispatch path for fast_topk_v2 and fast_topk_v2_raw. Templated on
// (K, kRawOutput). K is the top-k value (512 for V4-Flash, 1024 for V4-Pro);
// kRawOutput=false folds the page-table gather, kRawOutput=true emits raw
// row-local indices. The set of input tensors is the same modulo
// (page_table, page_size), which the caller has already validated.
template <uint32_t K, bool kRawOutput>
static void launch_dispatch(const TopKParams& params, uint32_t batch_size,
uint32_t max_seq_len, cudaStream_t stream) {
// Helper: build a cudaLaunchConfig with optional PDL + cluster attributes.
// The attribute storage must outlive cudaLaunchKernelEx (cfg.attrs points
// into it), so it lives in each call site below as a stack local.
auto make_cfg = [&](dim3 grid, dim3 block, size_t smem,
cudaLaunchAttribute* attrs, bool enable_cluster,
bool enable_pdl) {
cudaLaunchConfig_t cfg{};
cfg.gridDim = grid;
cfg.blockDim = block;
cfg.dynamicSmemBytes = static_cast<unsigned>(smem);
cfg.stream = stream;
int n = 0;
if (enable_pdl) {
attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[n].val.programmaticStreamSerializationAllowed = 1;
++n;
}
if (enable_cluster) {
attrs[n].id = cudaLaunchAttributeClusterDimension;
attrs[n].val.clusterDim = {1, kClusterSize, 1};
++n;
}
cfg.numAttrs = n;
cfg.attrs = n ? attrs : nullptr;
return cfg;
};
auto check_launch = [](cudaError_t err) {
TORCH_CHECK(err == cudaSuccess,
"fast_topk_v2 launch failed: ", cudaGetErrorString(err));
};
constexpr size_t kS1 = kStage1SMEM<K>;
constexpr size_t kS2 = kStage2SMEM<K>;
if (max_seq_len <= Small<K>::kMax1PassLength) {
setup_kernel_smem_once<&topk_short_transform<K, kRawOutput>, kS2>();
cudaLaunchAttribute attrs[2];
auto cfg = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs,
/*cluster=*/false, /*pdl=*/true);
check_launch(cudaLaunchKernelEx(
&cfg, topk_short_transform<K, kRawOutput>, params));
} else if (batch_size <= kNumClusters) {
constexpr size_t kFusedSMEM = kS1 > kS2 ? kS1 : kS2;
setup_kernel_smem_once<&topk_fused_transform<K, kRawOutput>, kFusedSMEM>();
cudaLaunchAttribute attrs[2];
auto cfg = make_cfg(dim3(batch_size, kClusterSize), dim3(kBlockSize),
kFusedSMEM, attrs, /*cluster=*/true, /*pdl=*/true);
check_launch(cudaLaunchKernelEx(
&cfg, topk_fused_transform<K, kRawOutput>, params));
} else {
const auto num_clusters = std::min<uint32_t>(batch_size, kNumClusters);
setup_kernel_smem_once<&topk_combine_preprocess<K, kRawOutput>, kS1>();
cudaLaunchAttribute attrs1[2];
auto cfg1 = make_cfg(dim3(num_clusters, kClusterSize), dim3(kBlockSize),
kS1, attrs1, /*cluster=*/true, /*pdl=*/true);
check_launch(cudaLaunchKernelEx(
&cfg1, topk_combine_preprocess<K, kRawOutput>, params));
setup_kernel_smem_once<&topk_combine_transform<K, kRawOutput>, kS2>();
cudaLaunchAttribute attrs2[2];
auto cfg2 = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs2,
/*cluster=*/false, /*pdl=*/true);
check_launch(cudaLaunchKernelEx(
&cfg2, topk_combine_transform<K, kRawOutput>, params));
}
}
// Top-level K dispatcher: validate the runtime topk argument and route to
// the right template instantiation.
template <bool kRawOutput>
static void launch_dispatch_k(int64_t topk, const TopKParams& params,
uint32_t batch_size, uint32_t max_seq_len,
cudaStream_t stream) {
if (topk == 512) {
launch_dispatch<512, kRawOutput>(params, batch_size, max_seq_len, stream);
} else if (topk == 1024) {
launch_dispatch<1024, kRawOutput>(params, batch_size, max_seq_len, stream);
} else {
TORCH_CHECK(false,
"fast_topk_v2 supports topk in {512, 1024}, got ", topk);
}
}
} // namespace vllm::dsv4_topk
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
const torch::Tensor& page_table, torch::Tensor& page_indices,
int64_t page_size, const torch::Tensor& workspace,
const torch::Tensor& metadata, int64_t topk) {
using namespace vllm::dsv4_topk;
CHECK_CUDA(scores);
CHECK_CUDA(seq_lens);
CHECK_CUDA(page_table);
CHECK_CUDA(page_indices);
CHECK_CUDA(workspace);
CHECK_CUDA(metadata);
CHECK_DTYPE(scores, torch::kFloat32);
CHECK_DTYPE(seq_lens, torch::kInt32);
CHECK_DTYPE(page_table, torch::kInt32);
CHECK_DTYPE(page_indices, torch::kInt32);
CHECK_DTYPE(workspace, torch::kInt32);
CHECK_DTYPE(metadata, torch::kInt32);
TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1,
"scores must be 2D with last stride 1");
TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous());
TORCH_CHECK(page_table.dim() == 2 && page_table.stride(1) == 1,
"page_table must be 2D with last stride 1");
TORCH_CHECK(page_indices.dim() == 2 && page_indices.is_contiguous() &&
page_indices.size(1) == topk,
"page_indices must be (B, topk) contiguous");
// workspace size is K-independent (it stages cluster-path ties whose
// count is bounded by kMaxTies, not K), so this check uses any K.
TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
workspace.size(1) == Large<512>::kWorkspaceInts,
"workspace must be (B, kWorkspaceInts) with last stride 1");
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 &&
metadata.is_contiguous(),
"metadata must be (B + 1, 4) contiguous");
const auto batch_size = static_cast<uint32_t>(scores.size(0));
TORCH_CHECK(seq_lens.size(0) == batch_size);
TORCH_CHECK(page_table.size(0) == batch_size);
TORCH_CHECK(page_indices.size(0) == batch_size);
TORCH_CHECK(workspace.size(0) == batch_size);
TORCH_CHECK(metadata.size(0) == batch_size + 1);
const auto max_seq_len = static_cast<uint32_t>(scores.size(1));
TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0,
"page_size must be a positive power of 2");
TORCH_CHECK(scores.stride(0) % 4 == 0,
"score stride must be a multiple of 4 (TMA 16-byte alignment)");
// page_bits = log2(page_size). __builtin_ctzll is a host-side compiler
// builtin available under C++17 (vLLM compiles host code with C++17).
const auto page_bits = static_cast<uint32_t>(
__builtin_ctzll(static_cast<unsigned long long>(page_size)));
TopKParams params{
.seq_lens =
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
.scores = scores.data_ptr<float>(),
.page_table = page_table.data_ptr<int32_t>(),
.page_indices = page_indices.data_ptr<int32_t>(),
.score_stride = scores.stride(0),
.page_table_stride = page_table.stride(0),
.workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
.metadata =
reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
.workspace_stride =
workspace.stride(0) * static_cast<int64_t>(sizeof(int32_t)),
.batch_size = batch_size,
.page_bits = page_bits,
};
launch_dispatch_k<false>(topk, params, batch_size, max_seq_len,
at::cuda::getCurrentCUDAStream().stream());
}
// Top-k only: skip the page-table gather and emit raw row-local indices.
// Same selection algorithm as fast_topk_v2; just doesn't touch a page
// table. Output semantics match torch.ops._C.persistent_topk and the V4
// indexer's existing topk_indices_buffer contract.
void fast_topk_v2_raw(const torch::Tensor& scores,
const torch::Tensor& seq_lens,
torch::Tensor& topk_indices,
const torch::Tensor& workspace,
const torch::Tensor& metadata,
int64_t topk) {
using namespace vllm::dsv4_topk;
CHECK_CUDA(scores);
CHECK_CUDA(seq_lens);
CHECK_CUDA(topk_indices);
CHECK_CUDA(workspace);
CHECK_CUDA(metadata);
CHECK_DTYPE(scores, torch::kFloat32);
CHECK_DTYPE(seq_lens, torch::kInt32);
CHECK_DTYPE(topk_indices, torch::kInt32);
CHECK_DTYPE(workspace, torch::kInt32);
CHECK_DTYPE(metadata, torch::kInt32);
TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1,
"scores must be 2D with last stride 1");
TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous());
TORCH_CHECK(topk_indices.dim() == 2 && topk_indices.is_contiguous() &&
topk_indices.size(1) == topk,
"topk_indices must be (B, topk) contiguous");
TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
workspace.size(1) == Large<512>::kWorkspaceInts,
"workspace must be (B, kWorkspaceInts) with last stride 1");
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 &&
metadata.is_contiguous(),
"metadata must be (B + 1, 4) contiguous");
const auto batch_size = static_cast<uint32_t>(scores.size(0));
TORCH_CHECK(seq_lens.size(0) == batch_size);
TORCH_CHECK(topk_indices.size(0) == batch_size);
TORCH_CHECK(workspace.size(0) == batch_size);
TORCH_CHECK(metadata.size(0) == batch_size + 1);
const auto max_seq_len = static_cast<uint32_t>(scores.size(1));
TORCH_CHECK(scores.stride(0) % 4 == 0,
"score stride must be a multiple of 4 (TMA 16-byte alignment)");
// page_table / page_bits are unused on the raw path; passing nullptr/0 is
// safe because every kernel call site is gated by `if constexpr
// (kRawOutput)` so the page-table loads are eliminated at compile time.
TopKParams params{
.seq_lens =
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
.scores = scores.data_ptr<float>(),
.page_table = nullptr,
.page_indices = topk_indices.data_ptr<int32_t>(),
.score_stride = scores.stride(0),
.page_table_stride = 0,
.workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
.metadata =
reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
.workspace_stride =
workspace.stride(0) * static_cast<int64_t>(sizeof(int32_t)),
.batch_size = batch_size,
.page_bits = 0,
};
launch_dispatch_k<true>(topk, params, batch_size, max_seq_len,
at::cuda::getCurrentCUDAStream().stream());
}
int64_t fast_topk_v2_workspace_ints() {
// Workspace size is K-independent (kMaxTies, not K, drives it).
return static_cast<int64_t>(vllm::dsv4_topk::Large<512>::kWorkspaceInts);
}
// Register impls here (instead of in torch_bindings.cpp) so they only exist
// when CMake compiles this source — i.e., when the target build has a
// compatible Hopper / Blackwell-datacenter arch. On other configs the schema
// remains defined but a call surfaces a clear "no impl" runtime error.
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("fast_topk_v2_plan", &fast_topk_v2_plan);
m.impl("fast_topk_v2", &fast_topk_v2);
m.impl("fast_topk_v2_raw", &fast_topk_v2_raw);
}
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) {
m.impl("fast_topk_v2_workspace_ints", &fast_topk_v2_workspace_ints);
}
+266
View File
@@ -0,0 +1,266 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Cluster top-k strategy for very large N. Uses Hopper thread-block clusters
// (cooperative_groups::this_cluster) to parallelize histogram + scatter across
// up to ``kClusterSize`` blocks per row. Each row is processed in two stages:
// stage 1: per-block histogram, all-reduce across the cluster, threshold
// scatter, and an epilogue that page-translates strictly-above
// entries to global memory and stages ties into a per-row workspace.
// stage 2: tie-break across the cluster's combined ties (run by cluster
// rank 0 in the fused kernel, or as a separate launch otherwise).
// Ported from
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh.
#pragma once
#include "common.cuh"
#include "ptx.cuh"
#include "utils.cuh"
#include <cooperative_groups.h>
#include <cstdint>
namespace vllm::dsv4_topk {
template <uint32_t K>
struct ClusterTopK {
static constexpr uint32_t kClusterSize = 8;
static constexpr uint32_t kHistBits = 10;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kElemPerStage = 8;
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
static constexpr uint32_t kNumStages = 4;
static constexpr uint32_t kMaxLength = kClusterSize * kNumStages * kSizePerStage;
static constexpr uint32_t kAboveBits = 11;
struct Smem {
uint64_t barrier[kNumStages];
uint32_t local_above_equal[kClusterSize];
uint32_t prefix_above_equal;
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
alignas(128) MatchBin match;
alignas(128) uint32_t warp_sum[kNumWarps];
uint32_t histogram[kHistBins];
alignas(128) float score_buffer[kNumStages][kSizePerStage];
Tie tie_buffer[kMaxTies];
};
// Per-row metadata produced by the plan kernel and consumed by the fused /
// stage-1 kernels. {batch_id, seq_len, has_next} arranged in an int4-sized
// 16-byte struct so the planner can do contiguous int32x4 stores.
struct alignas(16) Metadata {
uint32_t batch_id;
uint32_t seq_len;
bool has_next;
};
// Per-row workspace storing {(num_above, num_ties)} + the gathered ties.
struct WorkSpace {
uint2 metadata;
Tie ties[kMaxTies];
};
static constexpr uint32_t kWorkspaceInts = sizeof(WorkSpace) / sizeof(uint32_t);
VLLM_DSV4_DEVICE static void stage1_init(void* _smem) {
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
const auto smem = static_cast<Smem*>(_smem);
if (tx < kHistBins) smem->histogram[tx] = 0;
if (tx < kNumStages) ptx::mbarrier_init(&smem->barrier[tx], 1);
__syncthreads();
}
VLLM_DSV4_DEVICE static void stage1_prologue(const float* scores,
uint32_t length, void* _smem) {
if (threadIdx.x == 0) {
const auto smem = static_cast<Smem*>(_smem);
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
const auto length_aligned = (length + 3u) & ~3u;
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
if (stage >= num_stages) break;
const auto offset = stage * kSizePerStage;
const auto size = min(kSizePerStage, length_aligned - offset);
const auto size_bytes = size * sizeof(float);
const auto bar = &smem->barrier[stage];
ptx::tma_load(smem->score_buffer[stage], scores + offset, size_bytes,
bar);
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
}
}
}
VLLM_DSV4_DEVICE static void stage1(int32_t* indices, uint32_t length,
void* _smem, bool reuse = false) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// Local histogram.
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
const auto offset = stage * kSizePerStage;
if (offset >= length) break;
const auto size = min(kSizePerStage, length - offset);
if (lane_id == 0) ptx::mbarrier_wait(&smem->barrier[stage], 0);
__syncwarp();
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; ++i) {
const auto idx = tx + i * kBlockSize;
if (idx >= size) break;
const auto score = smem->score_buffer[stage][idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
atomicAdd(&smem->histogram[bin], 1);
}
}
static_assert(kHistBins <= kBlockSize);
// Two-shot all-reduce across the cluster.
{
auto cluster = cooperative_groups::this_cluster();
cluster.sync();
const auto cluster_rank = blockIdx.y;
const auto kLocalSize = kHistBins / kClusterSize;
const auto offset = kLocalSize * cluster_rank;
const auto src_tx = tx / kClusterSize;
const auto src_rank = tx % kClusterSize;
if (tx < kHistBins) {
const auto addr = &smem->histogram[offset + src_tx];
const auto src_addr = cluster.map_shared_rank(addr, src_rank);
*src_addr = warp_reduce_sum<kClusterSize>(*src_addr);
}
cluster.sync();
}
// Each block now holds the full cluster histogram. Find the threshold.
{
const auto value = tx < kHistBins ? smem->histogram[tx] : 0;
const auto warp_inc = warp_inclusive_sum(lane_id, value);
if (lane_id == kWarpThreads - 1) {
smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
const auto total_length = warp_reduce_sum(tmp);
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_inc;
const auto above = total_length - prefix_sum;
if (tx < kHistBins && above < K && above + value >= K) {
smem->counter_gt = smem->counter_eq = 0;
smem->match = {
.bin = tx,
.above_count = above,
.equal_count = value,
};
}
__syncthreads();
}
const auto thr_bin = smem->match.bin;
// Scatter strictly-above entries to `indices`, stash ties in tie_buffer.
#pragma unroll
for (uint32_t stage = 0; stage < kNumStages; stage++) {
const auto offset = stage * kSizePerStage;
if (offset >= length) break;
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; ++i) {
const auto buf_idx = tx + i * kBlockSize;
const auto global_idx = offset + buf_idx;
if (global_idx >= length) break;
const auto score = smem->score_buffer[stage][buf_idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
if (bin > thr_bin) {
indices[atomicAdd(&smem->counter_gt, 1)] = global_idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
}
}
}
if (reuse) {
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
if (tx < kHistBins) smem->histogram[tx] = 0;
if (tx < num_stages) ptx::mbarrier_arrive(&smem->barrier[tx]);
}
__syncthreads();
}
template <typename TParams>
VLLM_DSV4_DEVICE static void stage1_epilogue(TParams params,
uint32_t offset, void* _ws,
void* _smem) {
auto cluster = cooperative_groups::this_cluster();
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto local_above = smem->counter_gt;
const auto local_equal = smem->counter_eq;
const auto cluster_rank = blockIdx.y;
constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
static_assert(kAboveMask >= K);
static_assert(kMaxTies <= kBlockSize);
const auto idx_above = tx < local_above ? params.indices_in[tx] : 0;
const auto tie_value = tx < local_equal ? smem->tie_buffer[tx] : Tie{0, 0.0f};
// Push counts to remote shared memory to reduce inter-block latency.
if (tx < kClusterSize) {
const auto value = (local_equal << kAboveBits) | local_above;
const auto dst_addr = cluster.map_shared_rank(smem->local_above_equal, tx);
dst_addr[cluster_rank] = value;
}
// After this final sync, every block can read only its own smem (peer
// ranks may have already exited), so we don't touch remote smem again.
cluster.sync();
if (tx < kClusterSize) {
const auto value = tx < cluster_rank ? smem->local_above_equal[tx] : 0;
const auto kActiveMask = (1u << kClusterSize) - 1;
smem->prefix_above_equal = warp_reduce_sum<kClusterSize>(value, kActiveMask);
}
__syncthreads();
const auto prefix_packed = smem->prefix_above_equal;
const auto prefix_above = prefix_packed & kAboveMask;
const auto prefix_equal = prefix_packed >> kAboveBits;
// Page-translate strictly-above entries.
if (tx < local_above) {
params.write(tx + prefix_above, idx_above + offset);
}
// Stage ties into the per-row workspace (regular global writes).
const auto ws = static_cast<WorkSpace*>(_ws);
if (tx < local_equal && tx + prefix_equal < kMaxTies) {
ws->ties[tx + prefix_equal] = {tie_value.idx + offset, tie_value.score};
}
// Last cluster rank publishes the sums into ws->metadata.
if (cluster_rank == kClusterSize - 1 && tx == 0) {
const auto sum_above = prefix_above + local_above;
const auto sum_equal = prefix_equal + local_equal;
ws->metadata = make_uint2(sum_above, sum_equal);
}
}
template <typename TParams>
VLLM_DSV4_DEVICE static void transform(TParams params, const void* _ws,
void* _smem) {
const auto ws = static_cast<const WorkSpace*>(_ws);
const auto meta = &ws->metadata;
const auto num_above = meta->x;
const auto num_equal = meta->y;
if (num_above >= K || num_equal == 0) return;
const auto clamped_ties = min(num_equal, kMaxTies);
tie_handle_transform(ws->ties, clamped_ties, num_above, K, params, _smem);
}
};
} // namespace vllm::dsv4_topk
+219
View File
@@ -0,0 +1,219 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Shared types/utilities for the three DeepSeek V4 top-k strategies
// (Register / Streaming / Cluster). Ported from sglang's
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh.
#pragma once
#include "utils.cuh"
#include <cuda_fp16.h>
#include <cstdint>
namespace vllm::dsv4_topk {
inline constexpr uint32_t kMaxTopK = 1024;
inline constexpr uint32_t kBlockSize = 1024;
inline constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
// 1 element per thread in the tie-breaking pass.
inline constexpr uint32_t kMaxTies = 1024;
inline constexpr uint32_t kRadixBins = 256;
static_assert(kMaxTopK <= kBlockSize && kMaxTies <= kBlockSize);
// Always vectorize global loads as float4.
using Vec4 = AlignedVector<float, 4>;
// page_to_indices: convert a flat compressed-token index into a (block * page_size + offset)
// page-table-resolved index. page_size must be a power of 2; page_bits = log2(page_size).
VLLM_DSV4_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table,
uint32_t i, uint32_t page_bits) {
const uint32_t mask = (1u << page_bits) - 1u;
return (page_table[i >> page_bits] << page_bits) | (i & mask);
}
// Output-side description of how each strategy commits its top-k output.
//
// Two modes, picked at compile time via ``kRawOutput``:
// - kRawOutput=false (paged): fold the page-table gather into the output
// store. ``write(dst, src)`` emits ``page_to_indices(table, src, bits)``;
// ``transform(idx)`` reads ``indices_in[idx]`` and re-emits via the
// page lookup. This is the original kernel behavior.
// - kRawOutput=true (raw): skip the page lookup entirely. The kernel
// just writes row-local raw indices, matching ``persistent_topk``'s
// output contract. ``page_table`` and ``page_bits`` are unused; the
// compiler eliminates the dead loads via ``if constexpr``.
template <bool kRawOutput>
struct TransformParamsT {
const int32_t* __restrict__ page_table;
const int32_t* __restrict__ indices_in;
int32_t* __restrict__ indices_out;
uint32_t page_bits;
VLLM_DSV4_DEVICE void transform(uint32_t idx) const {
if constexpr (kRawOutput) {
indices_out[idx] = static_cast<int32_t>(indices_in[idx]);
} else {
indices_out[idx] =
page_to_indices(page_table, indices_in[idx], page_bits);
}
}
VLLM_DSV4_DEVICE void write(uint32_t dst, uint32_t src) const {
if constexpr (kRawOutput) {
indices_out[dst] = static_cast<int32_t>(src);
} else {
indices_out[dst] = page_to_indices(page_table, src, page_bits);
}
}
};
// Back-compat alias. The four kernels in fast_topk_v2.cu instantiate both
// variants explicitly via templates.
using TransformParams = TransformParamsT<false>;
struct alignas(16) MatchBin {
uint32_t bin;
uint32_t above_count;
uint32_t equal_count;
};
struct alignas(8) Tie {
uint32_t idx;
float score;
};
// Shared-memory layout for the final tie-breaking radix pass. Reused by both
// the streaming kernel (overlapping `score_buffer`) and the cluster kernel.
struct TieHandleSmem {
alignas(128) uint32_t counter;
alignas(128) MatchBin match;
uint32_t histogram[kRadixBins];
uint32_t warp_sum[kNumWarps];
};
// Order-preserving fp32 -> uint key, truncated to the top kBits. Used for the
// coarse histogram pass.
template <uint32_t kBits>
VLLM_DSV4_DEVICE uint32_t extract_coarse_bin(float x) {
static_assert(0 < kBits && kBits < 15);
__half h = __float2half_rn(x);
uint16_t bits = __half_as_ushort(h);
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
: static_cast<uint16_t>(bits | 0x8000);
return key >> (16 - kBits);
}
// Full 32-bit order-preserving key, used in tie-breaking.
VLLM_DSV4_DEVICE uint32_t extract_exact_bin(float x) {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
VLLM_DSV4_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
static_assert(kWarpThreads == 32);
#pragma unroll
for (uint32_t offset = 1; offset < 32; offset *= 2) {
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
if (lane_id >= offset) val += n;
}
return val;
}
// Fast path when seq_len <= K: identity mapping, padded to K with -1.
template <typename TParams>
VLLM_DSV4_DEVICE void trivial_transform(const TParams& params, uint32_t length,
uint32_t K) {
const auto tx = threadIdx.x;
if (tx < length) {
params.write(tx, tx);
} else if (tx < K) {
params.indices_out[tx] = -1;
}
}
// Tie-break the threshold-bin candidates that didn't fit in the strict-above
// region. One block-wide radix pass over the full 32-bit key (fp32 bit
// pattern, with idx as a secondary key). Writes at most `K - num_above`
// entries via params.write(...).
template <typename TParams>
VLLM_DSV4_DEVICE void tie_handle_transform(const Tie* __restrict__ ties,
uint32_t num_ties, uint32_t num_above,
uint32_t K, TParams params,
void* _smem) {
auto* smem = static_cast<TieHandleSmem*>(_smem);
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
const bool has_elem = tx < num_ties;
const auto tie = has_elem ? ties[tx] : Tie{0, 0.0f};
const uint32_t key = extract_exact_bin(tie.score);
const uint32_t idx = tie.idx;
bool active = has_elem;
uint32_t topk_remain = K - num_above;
uint32_t write_pos = K;
smem->counter = 0;
__syncthreads();
// 256 bins / 32 lanes = 8 warps span the histogram inter-warp prefix.
constexpr uint32_t kRadixWarps = kRadixBins / kWarpThreads;
#pragma unroll
for (int round = 0; round < 4; round++) {
const uint32_t shift = 24 - round * 8;
const uint32_t bin = (key >> shift) & 0xFFu;
// 1. Histogram.
if (tx < kRadixBins) smem->histogram[tx] = 0;
__syncthreads();
if (active) atomicAdd(&smem->histogram[bin], 1);
__syncthreads();
// 2. Two-pass prefix sum across the 256 bins.
uint32_t hist_val = 0;
uint32_t warp_inc = 0;
if (tx < kRadixBins) {
hist_val = smem->histogram[tx];
warp_inc = warp_inclusive_sum(lane_id, hist_val);
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
if (tx < kRadixBins) {
const auto tmp = (lane_id < kRadixWarps) ? smem->warp_sum[lane_id] : 0;
const auto total = warp_reduce_sum(tmp);
const auto inter = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
const auto prefix = inter + warp_inc;
const auto above = total - prefix;
// 3. Find threshold bin.
if (above < topk_remain && above + hist_val >= topk_remain) {
smem->match = {tx, above, topk_remain - above};
}
}
__syncthreads();
const auto thr = smem->match.bin;
const auto n_above = smem->match.above_count;
// 4. Scatter.
if (active) {
if (bin > thr) {
write_pos = num_above + atomicAdd(&smem->counter, 1);
active = false;
} else if (bin < thr) {
active = false;
} else if (round == 3) {
write_pos = K - atomicAdd(&smem->match.equal_count, -1u);
}
// bin == thr && round < 3: stay active for the next radix round.
}
topk_remain -= n_above;
if (topk_remain == 0) break;
}
if (write_pos < K) params.write(write_pos, idx);
}
} // namespace vllm::dsv4_topk
+66
View File
@@ -0,0 +1,66 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Thin wrappers around the CUDA PTX intrinsics used by the top-k pipeline.
// All of these require sm_90+. Ported from sglang's
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh.
#pragma once
#include "utils.cuh"
#include <cuda/ptx>
#include <cstdint>
namespace vllm::dsv4_topk::ptx {
VLLM_DSV4_DEVICE void mbarrier_init(uint64_t* addr, uint32_t arrives) {
cuda::ptx::mbarrier_init(addr, arrives);
}
VLLM_DSV4_DEVICE void mbarrier_arrive(uint64_t* addr) {
cuda::ptx::mbarrier_arrive(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta,
cuda::ptx::space_shared, addr);
}
VLLM_DSV4_DEVICE void mbarrier_arrive_expect_tx(uint64_t* addr, uint32_t tx) {
cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
cuda::ptx::scope_cta,
cuda::ptx::space_shared, addr, tx);
}
VLLM_DSV4_DEVICE void mbarrier_wait(uint64_t* addr, uint32_t phase) {
while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
cuda::ptx::scope_cta, addr,
phase))
;
}
VLLM_DSV4_DEVICE void tma_load(void* dst, const void* src, uint32_t num_bytes,
uint64_t* mbar) {
cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global,
dst, src, num_bytes, mbar);
}
// elect.sync: pick a single arbitrary thread out of an active mask. Used to
// fire a single TMA load per warp without the full ``if (tx == 0)`` cost.
VLLM_DSV4_DEVICE uint32_t elect_sync() {
uint32_t pred = 0;
asm volatile(
"{\n\t"
".reg .pred %%px;\n\t"
"elect.sync _|%%px, %1;\n\t"
"@%%px mov.s32 %0, 1;\n\t"
"}"
: "+r"(pred)
: "r"(0xFFFFFFFF));
return pred;
}
VLLM_DSV4_DEVICE bool elect_sync_cta(uint32_t tx) {
const auto warp_id = tx / 32;
const auto uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
return (uniform_warp_id == 0 && elect_sync());
}
} // namespace vllm::dsv4_topk::ptx
+314
View File
@@ -0,0 +1,314 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Register-resident top-k strategy for the DeepSeek V4 indexer (small N
// fast path). One block per row; up to ``kMax2PassLength`` scores per row
// streamed through registers, with a single 12-bit-coarse radix pass and
// a final tie-break round. Ported from
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh.
#pragma once
#include "common.cuh"
#include "ptx.cuh"
#include "utils.cuh"
#include <cfloat>
#include <cstdint>
namespace vllm::dsv4_topk {
template <uint32_t K>
struct RegisterTopK {
static constexpr uint32_t kHistBits = 12;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kVecsPerThread = 4;
static constexpr uint32_t kMaxTolerance = 0;
// Length covered by registers in a single pass.
static constexpr uint32_t kMax1PassLength = kVecsPerThread * 4 * kBlockSize;
// Extra length staged through shared memory in the 2-pass path.
static constexpr uint32_t kMaxExtraLength = kMax1PassLength;
static constexpr uint32_t kMax2PassLength = kMax1PassLength + kMaxExtraLength;
struct Smem {
using HistVec = AlignedVector<uint32_t, kHistBins / kBlockSize>;
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
uint64_t mbarrier; // for the cp.async.bulk in the 2-pass path
MatchBin match;
uint32_t warp_sum[kNumWarps];
union {
uint32_t histogram[kHistBins];
HistVec histogram_vec[kBlockSize];
Tie tie_buffer[kMaxTies];
};
alignas(16) float score_buffer[kMaxExtraLength];
};
template <bool kIs2Pass = false>
VLLM_DSV4_DEVICE static void run(const float* scores, int32_t* indices,
uint32_t length, void* _smem,
bool use_pdl = false) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
// Init histogram + counters.
{
typename Smem::HistVec hist_vec;
hist_vec.fill(0);
smem->histogram_vec[tx] = hist_vec;
if (tx == 0) {
smem->counter_gt = smem->counter_eq = 0;
if constexpr (kIs2Pass) {
ptx::mbarrier_init(&smem->mbarrier, 1);
}
}
__syncthreads();
}
if (use_pdl) pdl_wait_primary<true>();
// Stream the first `kMax1PassLength` scores into registers.
Vec4 local[kVecsPerThread];
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
const uint32_t base = (tx + v * kBlockSize) * 4;
if (base >= length) break;
local[v].load(scores, tx + v * kBlockSize);
}
// Issue the 2-pass TMA prefetch (next chunk of scores into smem).
if constexpr (kIs2Pass) {
if (ptx::elect_sync_cta(tx)) {
const auto length_aligned = (length + 3u - kMax1PassLength) & ~3u;
const auto size_bytes = length_aligned * sizeof(float);
ptx::tma_load(smem->score_buffer, scores + kMax1PassLength, size_bytes,
&smem->mbarrier);
ptx::mbarrier_arrive_expect_tx(&smem->mbarrier, size_bytes);
}
__syncwarp();
}
// Phase 1: histogram via shared-memory atomics.
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
if constexpr (!kIs2Pass) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if (idx >= length) goto LABEL_ACC_FINISH;
}
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local[v][e])],
1);
}
}
if constexpr (kIs2Pass) {
if (lane_id == 0) ptx::mbarrier_wait(&smem->mbarrier, 0);
__syncwarp();
for (uint32_t i = tx; i + kMax1PassLength < length; i += kBlockSize) {
const auto val = smem->score_buffer[i];
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(val)], 1);
}
}
[[maybe_unused]] LABEL_ACC_FINISH:
__syncthreads();
// Phase 2: prefix scan over the histogram, locate the threshold bin.
{
constexpr uint32_t kItems = kHistBins / kBlockSize;
uint32_t orig[kItems];
const auto hist_vec = smem->histogram_vec[tx];
uint32_t tmp_local_sum = 0;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
orig[i] = hist_vec[i];
tmp_local_sum += orig[i];
}
const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
const auto warp_exc = warp_inc - tmp_local_sum;
if (lane_id == kWarpThreads - 1) {
smem->warp_sum[warp_id] = warp_inc;
}
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
// Exactly one bin satisfies above < K && above + count >= K.
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_exc;
#pragma unroll
for (uint32_t i = 0; i < kItems; ++i) {
prefix_sum += orig[i];
const auto above = length - prefix_sum;
if (above < K && above + orig[i] >= K) {
smem->match = {
.bin = tx * kItems + i,
.above_count = above,
.equal_count = orig[i],
};
}
}
__syncthreads();
}
const auto thr_bin = smem->match.bin;
const auto num_above = smem->match.above_count;
const auto num_equal = smem->match.equal_count;
// Phase 3: Scatter.
// - bin > thr -> write directly to output (strictly above).
// - bin == thr -> when no tie-break is needed, admit first-come;
// otherwise stash into tie_buffer for phase 4.
const bool need_tiebreak = (num_equal + num_above > K + kMaxTolerance);
const auto topk_indices = indices;
const auto tie_buffer = smem->tie_buffer;
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if constexpr (!kIs2Pass) {
if (idx >= length) goto LABEL_SCATTER_DONE;
}
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
if (bin > thr_bin) {
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (need_tiebreak) {
if (pos < kMaxTies) {
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
}
} else {
if (const auto which = pos + num_above; which < K) {
topk_indices[which] = idx;
}
}
}
}
// 2-pass: pull the next chunk in from the staged smem buffer.
if constexpr (kIs2Pass) {
local[v].load(smem->score_buffer, tx + v * kBlockSize);
}
}
if constexpr (kIs2Pass) {
#pragma unroll
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
#pragma unroll
for (uint32_t e = 0; e < 4; ++e) {
const uint32_t idx =
(tx + v * kBlockSize) * 4 + e + kMax1PassLength;
if (idx >= length) goto LABEL_SCATTER_DONE;
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
if (bin > thr_bin) {
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (need_tiebreak) {
if (pos < kMaxTies) {
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
}
} else {
if (const auto which = pos + num_above; which < K) {
topk_indices[which] = idx;
}
}
}
}
}
}
[[maybe_unused]] LABEL_SCATTER_DONE:
if (!need_tiebreak) return;
// Phase 4: tie-break within the threshold bin. We assume num_ties <=
// kBlockSize (one block of ties), so each thread takes one tied element,
// counts the number of tied elements with strictly higher (score, -idx),
// and writes to output if its rank is below the remaining quota.
__syncthreads();
static_assert(kMaxTies <= kBlockSize);
const uint32_t num_ties = min(num_equal, kMaxTies);
const uint32_t topk_remain = K - num_above;
const auto is_greater = [](const Tie& a, const Tie& b) {
return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
};
if (num_ties <= kWarpThreads) {
static_assert(kWarpThreads <= kNumWarps);
if (lane_id >= num_ties || warp_id >= num_ties) return;
const uint32_t mask = (1ull << num_ties) - 1u;
const auto tie = tie_buffer[lane_id];
const auto target_tie = tie_buffer[warp_id];
const bool pred = is_greater(tie, target_tie);
const auto rank =
static_cast<uint32_t>(__popc(__ballot_sync(mask, pred)));
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target_tie.idx;
}
} else if (num_ties <= kWarpThreads * 2) {
// 64x64 case: each thread takes 2 elements.
const auto lane_id_1 = lane_id + kWarpThreads;
const auto warp_id_1 = warp_id + kWarpThreads;
const auto invalid = Tie{.idx = 0xFFFFFFFFu, .score = -FLT_MAX};
const auto tie_0 = tie_buffer[lane_id];
const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
{
const auto target = tie_buffer[warp_id];
const bool pred_0 = is_greater(tie_0, target);
const bool pred_1 = is_greater(tie_1, target);
const auto rank_0 =
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
const auto rank_1 =
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target.idx;
}
}
if (warp_id_1 < num_ties) {
const auto target = tie_buffer[warp_id_1];
const bool pred_0 = is_greater(tie_0, target);
const bool pred_1 = is_greater(tie_1, target);
const auto rank_0 =
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
const auto rank_1 =
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
const auto rank = rank_0 + rank_1;
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target.idx;
}
}
} else {
[[unlikely]];
// Block-wide fallback. Rarely reached.
for (auto i = warp_id; i < num_ties; i += kNumWarps) {
const auto target_tie = tie_buffer[i];
uint32_t local_rank = 0;
for (auto j = lane_id; j < num_ties; j += kWarpThreads) {
const auto tie = tie_buffer[j];
if (is_greater(tie, target_tie)) local_rank++;
}
const auto rank = warp_reduce_sum(local_rank);
if (lane_id == 0 && rank < topk_remain) {
topk_indices[num_above + rank] = target_tie.idx;
}
}
}
}
template <typename TParams>
VLLM_DSV4_DEVICE static void transform(TParams params) {
__syncthreads();
if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
}
};
} // namespace vllm::dsv4_topk
+209
View File
@@ -0,0 +1,209 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Streaming top-k strategy for medium N. Uses a TMA-driven double-buffered
// histogram pass + scatter pass over chunks of `kSizePerStage` floats.
// Ported from
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh.
#pragma once
#include "common.cuh"
#include "ptx.cuh"
#include "utils.cuh"
#include <cfloat>
#include <cstdint>
namespace vllm::dsv4_topk {
template <uint32_t K>
struct StreamingTopK {
static constexpr uint32_t kHistBits = 12;
static constexpr uint32_t kHistBins = 1 << kHistBits;
static constexpr uint32_t kElemPerStage = 8;
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
static constexpr uint32_t kNumStages = 2; // double buffer
static constexpr uint32_t kHistItems = kHistBins / kBlockSize; // 4
static_assert(kHistItems * kBlockSize == kHistBins);
using HistVec = AlignedVector<uint32_t, kHistItems>;
struct Smem {
// [phase = 0 (histogram) | 1 (scatter)] x [buffer = 0 | 1]
uint64_t barrier[2][kNumStages];
alignas(128) uint32_t counter_gt;
alignas(128) uint32_t counter_eq;
alignas(128) MatchBin match;
alignas(128) uint32_t warp_sum[kNumWarps];
union {
uint32_t histogram[kHistBins];
HistVec histogram_vec[kBlockSize];
Tie tie_buffer[kMaxTies];
};
union {
float score_buffer[kNumStages][kSizePerStage];
TieHandleSmem stage2; // reused for the tie-handling phase
};
};
// length must be 4-aligned (caller rounds up); TMA wants 16-byte alignment.
template <bool kIsScatter>
VLLM_DSV4_DEVICE static void issue_tma(const float* scores, uint32_t stage,
uint32_t length, Smem* smem) {
const auto buf_idx = stage % kNumStages;
const auto offset = stage * kSizePerStage;
const auto size = min(kSizePerStage, length - offset);
const auto size_bytes = size * sizeof(float);
const auto bar = &smem->barrier[kIsScatter][buf_idx];
ptx::tma_load(smem->score_buffer[buf_idx], scores + offset, size_bytes,
bar);
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
}
// Unified streaming pass. kIsScatter=false: build histogram (phase A).
// kIsScatter=true: scatter using the threshold bin (phase C). Each barrier
// is reused across iterations via the reuse-arrive pattern.
template <bool kIsScatter>
VLLM_DSV4_DEVICE static void stream_pass(const float* scores, uint32_t length,
uint32_t thr_bin,
int32_t* s_topk_indices,
Smem* smem) {
const auto tx = threadIdx.x;
const auto num_iters = (length + kSizePerStage - 1) / kSizePerStage;
const auto lane_id = tx % kWarpThreads;
const auto length_aligned = (length + 3u) & ~3u;
if (tx == 0) {
#pragma unroll
for (uint32_t i = 0; i < kNumStages; i++) {
if (i >= num_iters) break;
issue_tma<kIsScatter>(scores, i, length_aligned, smem);
}
}
for (uint32_t iter = 0; iter < num_iters; iter++) {
const auto buf_idx = iter % kNumStages;
const auto offset = iter * kSizePerStage;
const auto this_size = min(kSizePerStage, length - offset);
if (lane_id == 1) {
const auto phase_bit = (iter / kNumStages) & 1;
ptx::mbarrier_wait(&smem->barrier[kIsScatter][buf_idx], phase_bit);
}
__syncwarp();
#pragma unroll
for (uint32_t i = 0; i < kElemPerStage; i++) {
const auto local_idx = tx + i * kBlockSize;
if (local_idx >= this_size) break;
const auto score = smem->score_buffer[buf_idx][local_idx];
const auto bin = extract_coarse_bin<kHistBits>(score);
if constexpr (kIsScatter) {
const auto global_idx = offset + local_idx;
if (bin > thr_bin) {
const auto pos = atomicAdd(&smem->counter_gt, 1);
if (pos < K) s_topk_indices[pos] = global_idx;
} else if (bin == thr_bin) {
const auto pos = atomicAdd(&smem->counter_eq, 1);
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
}
} else {
atomicAdd(&smem->histogram[bin], 1);
}
}
__syncthreads();
if (tx == 0) {
if (const auto next_iter = iter + kNumStages; next_iter < num_iters) {
issue_tma<kIsScatter>(scores, next_iter, length_aligned, smem);
}
}
}
}
// Phase B: locate threshold bin via warp-level prefix scan.
VLLM_DSV4_DEVICE static void find_threshold(uint32_t length, Smem* smem) {
const auto tx = threadIdx.x;
const auto lane_id = tx % kWarpThreads;
const auto warp_id = tx / kWarpThreads;
uint32_t orig[kHistItems];
const auto hist_vec = smem->histogram_vec[tx];
uint32_t local_sum = 0;
#pragma unroll
for (uint32_t i = 0; i < kHistItems; ++i) {
orig[i] = hist_vec[i];
local_sum += orig[i];
}
const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
const auto warp_exc = warp_inc - local_sum;
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
__syncthreads();
const auto tmp = smem->warp_sum[lane_id];
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
prefix_sum += warp_exc;
#pragma unroll
for (uint32_t i = 0; i < kHistItems; ++i) {
prefix_sum += orig[i];
const auto above = length - prefix_sum;
if (above < K && above + orig[i] >= K) {
smem->match = {
.bin = tx * kHistItems + i,
.above_count = above,
.equal_count = orig[i],
};
}
}
__syncthreads();
}
VLLM_DSV4_DEVICE static void run(const float* scores, uint32_t length,
int32_t* topk_indices, void* _smem) {
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
__builtin_assume(tx < kBlockSize);
{
HistVec zero;
zero.fill(0);
smem->histogram_vec[tx] = zero;
if (tx < 2 * kNumStages) {
const auto base_barrier = &smem->barrier[0][0];
ptx::mbarrier_init(&base_barrier[tx], 1);
}
if (tx == 0) {
smem->counter_gt = 0;
smem->counter_eq = 0;
}
__syncthreads();
}
// Phase A: histogram.
stream_pass<false>(scores, length, 0, nullptr, smem);
// Phase B: threshold bin.
find_threshold(length, smem);
// Phase C: scatter.
stream_pass<true>(scores, length, smem->match.bin, topk_indices, smem);
}
template <typename TParams>
VLLM_DSV4_DEVICE static void transform(TParams params, void* _smem) {
// Phase D: page-translate above entries, then refine ties.
const auto smem = static_cast<Smem*>(_smem);
const auto tx = threadIdx.x;
const auto num_above = smem->match.above_count;
if (tx < num_above) params.transform(tx);
const auto num_equal = smem->counter_eq;
if (num_above >= K || num_equal == 0) return;
const auto clamped_ties = min(num_equal, kMaxTies);
tie_handle_transform(smem->tie_buffer, clamped_ties, num_above, K, params,
&smem->stage2);
}
};
} // namespace vllm::dsv4_topk
+75
View File
@@ -0,0 +1,75 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Minimal device-side utilities used by the DeepSeek V4 indexer top-k port.
// Replaces sgl_kernel/{utils,warp,vec,type}.cuh — we only need the bits the
// top-k kernels actually touch.
#pragma once
#include <cstddef>
#include <cstdint>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
namespace vllm::dsv4_topk {
#define VLLM_DSV4_DEVICE __forceinline__ __device__
inline constexpr uint32_t kWarpThreads = 32u;
inline constexpr uint32_t kFullMask = 0xffffffffu;
// Programmatic Dependent Launch (sm_90+). When enabled, the kernel waits for
// the predecessor on the same stream to advance past its dependents-launch
// trigger before doing anything memory-dependent. Used to overlap the
// fp8_paged_mqa_logits epilogue with the first stage of top-k.
template <bool kUsePDL>
VLLM_DSV4_DEVICE void pdl_wait_primary() {
if constexpr (kUsePDL) {
asm volatile("griddepcontrol.wait;" ::: "memory");
}
}
template <bool kUsePDL>
VLLM_DSV4_DEVICE void pdl_trigger_secondary() {
if constexpr (kUsePDL) {
asm volatile("griddepcontrol.launch_dependents;" :::);
}
}
// Warp-level XOR-shuffle reduce. kThreads must be a power of 2 and <= 32.
template <uint32_t kThreads = kWarpThreads, typename T>
VLLM_DSV4_DEVICE T warp_reduce_sum(T value, uint32_t active_mask = kFullMask) {
#pragma unroll
for (auto offset = kThreads >> 1; offset > 0; offset >>= 1) {
value = value + __shfl_xor_sync(active_mask, value, offset, 32);
}
return value;
}
// 128-bit-aligned vector of N elements of T (N must be a power of 2, total
// size <= 16 bytes). Used for vectorized loads/stores into shared memory.
template <typename T, std::size_t N>
struct alignas(sizeof(T) * N) AlignedVector {
static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two");
static_assert(sizeof(T) * N <= 16,
"AlignedVector exceeds the 128-bit CUDA vector limit");
T data[N];
VLLM_DSV4_DEVICE void load(const void* ptr, std::size_t offset = 0) {
*reinterpret_cast<AlignedVector*>(this) =
reinterpret_cast<const AlignedVector*>(ptr)[offset];
}
VLLM_DSV4_DEVICE void store(void* ptr, std::size_t offset = 0) const {
reinterpret_cast<AlignedVector*>(ptr)[offset] = *this;
}
VLLM_DSV4_DEVICE void fill(T value) {
#pragma unroll
for (std::size_t i = 0; i < N; ++i) data[i] = value;
}
VLLM_DSV4_DEVICE T& operator[](std::size_t i) { return data[i]; }
VLLM_DSV4_DEVICE const T& operator[](std::size_t i) const { return data[i]; }
};
} // namespace vllm::dsv4_topk
+3 -6
View File
@@ -137,18 +137,15 @@ fused_add_rms_norm_static_fp8_quant_kernel(
_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;
#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);
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
static_cast<float>(out_norm), scale_inv);
}
}
}
@@ -141,7 +141,7 @@ struct cutlass_3x_gemm_sm100 {
sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelSchedule>::CollectiveOp;
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
};
@@ -125,7 +125,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
MainloopScheduler
>::CollectiveOp>;
using KernelType = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using KernelType = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
struct GemmKernel : public KernelType {};
@@ -92,7 +92,7 @@ struct cutlass_3x_gemm_sm100_fp8 {
// -----------------------------------------------------------
// Kernel definition
// -----------------------------------------------------------
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
};
+34 -2
View File
@@ -125,6 +125,40 @@ 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);
// DeepSeek V4 indexer top-k (k = 512). Hopper (sm_90a) and Blackwell
// datacenter (sm_100/sm_103) — needs thread-block clusters, TMA, and PDL.
// Two-step API:
// 1. fast_topk_v2_plan inspects the seq_lens distribution and writes a
// cluster_threshold + per-row Metadata into a (B+1, 4) int32 tensor. The
// plan is amortized when cudagraph-captured: once per shape, reused across
// layers.
// 2. fast_topk_v2 selects the top-512 indices per row, folds the page-table
// gather into the radix store, and writes (B, 512) int32 page indices.
// Dispatches per row to one of three strategies (Register / Streaming /
// Cluster) using the planned threshold.
//
// Returns the size in int32s of the per-row workspace required by
// fast_topk_v2 (allocate `(B, fast_topk_v2_workspace_ints())` int32 contig).
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
int64_t static_cluster_threshold);
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
const torch::Tensor& page_table, torch::Tensor& page_indices,
int64_t page_size, const torch::Tensor& workspace,
const torch::Tensor& metadata, int64_t topk);
// Top-k only, no page-table fold-in. Same selection as fast_topk_v2 but
// emits raw row-local indices into ``topk_indices`` (drop-in for
// persistent_topk's output contract). topk must be one of {512, 1024}.
void fast_topk_v2_raw(const torch::Tensor& scores,
const torch::Tensor& seq_lens,
torch::Tensor& topk_indices,
const torch::Tensor& workspace,
const torch::Tensor& metadata,
int64_t topk);
int64_t fast_topk_v2_workspace_ints();
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& weight, torch::Tensor& scale,
double epsilon);
@@ -163,8 +197,6 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
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);
+34 -39
View File
@@ -7,23 +7,23 @@
namespace vllm {
template <typename scalar_t, typename cache_t, bool IS_NEOX>
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr,
const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
const bool inverse) {
int x_index, y_index;
float cos_f, sin_f;
if (IS_NEOX) {
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index));
sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index));
cos_f = VLLM_LDG(cos_ptr + x_index);
sin_f = VLLM_LDG(sin_ptr + x_index);
} else {
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2));
sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2));
cos_f = VLLM_LDG(cos_ptr + x_index / 2);
sin_f = VLLM_LDG(sin_ptr + x_index / 2);
}
if (inverse) {
sin_f = -sin_f;
@@ -34,7 +34,7 @@ inline __device__ void apply_token_rotary_embedding(
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
}
template <typename scalar_t, typename cache_t, bool IS_NEOX>
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_rotary_embedding(
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
@@ -43,14 +43,14 @@ inline __device__ void apply_rotary_embedding(
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const cache_t* cache_ptr, const int head_size, const int num_heads,
const float* 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 int embed_dim = rot_dim / 2;
const cache_t* cos_ptr = cache_ptr;
const cache_t* sin_ptr = cache_ptr + embed_dim;
const float* cos_ptr = cache_ptr;
const float* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
@@ -58,7 +58,7 @@ inline __device__ void apply_rotary_embedding(
const int64_t token_head =
token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
}
@@ -69,13 +69,13 @@ inline __device__ void apply_rotary_embedding(
const int64_t token_head =
token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
}
}
}
template <typename scalar_t, typename cache_t, bool IS_NEOX>
template <typename scalar_t, bool IS_NEOX>
__global__ void rotary_embedding_kernel(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
// [num_tokens]
@@ -86,15 +86,15 @@ __global__ void rotary_embedding_kernel(
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const cache_t* __restrict__ cos_sin_cache, // [max_position, rot_dim]
const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
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 token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const cache_t* cache_ptr = cos_sin_cache + pos * rot_dim;
const float* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
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);
@@ -168,28 +168,23 @@ 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", [&] {
using query_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(
cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
using cache_t = scalar_t;
if (is_neox) {
vllm::rotary_embedding_kernel<query_t, cache_t, true>
<<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
key.has_value() ? key->data_ptr<query_t>() : nullptr,
cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
key_stride, head_stride, num_heads, num_kv_heads, head_size,
rope_dim_offset, inverse);
} else {
vllm::rotary_embedding_kernel<query_t, cache_t, false>
<<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
key.has_value() ? key->data_ptr<query_t>() : nullptr,
cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
key_stride, head_stride, num_heads, num_kv_heads, head_size,
rope_dim_offset, inverse);
}
});
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);
} 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);
}
});
}
+4 -59
View File
@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
// Query occupancy for the instantiation that will actually launch;
// overestimating it deadlocks the cooperative barrier.
int occupancy = 1;
cudaError_t occ_err = cudaSuccess;
if (vec_size == 4) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
smem_size);
} else if (vec_size == 2) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
smem_size);
} else {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
smem_size);
}
TORCH_CHECK(occ_err == cudaSuccess,
"persistent_topk occupancy query failed: ",
cudaGetErrorString(occ_err));
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
smem_size);
if (occupancy < 1) occupancy = 1;
// The cooperative spin-wait barrier only runs when at least one row hits
// the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
// early-exit, so oversubscription can't deadlock and headroom is wasted.
const bool needs_cooperative =
static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
const uint32_t hw_resident_cap =
static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
uint32_t max_resident_ctas = hw_resident_cap;
if (needs_cooperative) {
// Reserve one CTA per SM when occupancy allows; fall back to a single
// CTA when occupancy == 1 (the most deadlock-prone case — any straggler
// kernel that takes the only slot on one SM hangs the barrier). Never
// drop below one full group's worth.
uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
if (max_resident_ctas >= headroom + ctas_per_group) {
max_resident_ctas -= headroom;
}
}
uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
static_cast<uint32_t>(num_rows));
if (num_groups == 0) num_groups = 1;
uint32_t total_ctas = num_groups * ctas_per_group;
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
// instead of deadlocking. Only relevant when needs_cooperative.
if (needs_cooperative && total_ctas > hw_resident_cap) {
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
"persistent_topk would oversubscribe and the FilteredTopK "
"fallback requires >=128KB smem per block (have ",
max_smem_per_block, "). total_ctas=", total_ctas,
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
", smem=", smem_size, ").");
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
stream);
TORCH_CHECK(status == cudaSuccess,
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
return;
}
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
"workspace too small, need ", state_bytes, " bytes");
+20 -6
View File
@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
// SwiGLU activation with input clamping.
ops.def(
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
"-> ()");
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -221,6 +215,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"Tensor workspace, int k, int max_seq_len) -> ()");
ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
// DeepSeek V4 indexer top-k (k=512), ported from sglang's topk_v2 family.
// Built for sm_90a (Hopper) + sm_100a/sm_103 (Blackwell datacenter).
// Schema only here; impl is registered in csrc/deepseek_v4/fast_topk_v2.cu
// so it's only present when CMake compiles the source for a supported arch.
ops.def(
"fast_topk_v2_plan(Tensor seq_lens, Tensor! metadata, "
"int static_cluster_threshold) -> ()");
ops.def(
"fast_topk_v2(Tensor scores, Tensor seq_lens, Tensor page_table, "
"Tensor! page_indices, int page_size, Tensor workspace, "
"Tensor metadata, int topk) -> ()");
ops.def(
"fast_topk_v2_raw(Tensor scores, Tensor seq_lens, "
"Tensor! topk_indices, Tensor workspace, Tensor metadata, int topk)"
" -> ()");
ops.def("fast_topk_v2_workspace_ints() -> int");
// Layernorm-quant
// Apply Root Mean Square (RMS) Normalization to the input tensor.
ops.def(
+5 -28
View File
@@ -538,11 +538,9 @@ 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 \
# numactl CLI for NUMA binding at runtime
numactl && \
libnuma-dev && \
# Fixes nccl_allocator requiring nccl.h at runtime
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -585,7 +583,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
ARG FLASHINFER_VERSION=0.6.8.post1
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
&& flashinfer show-config \
&& flashinfer download-cubin
# ============================================================
# OPENAI API SERVER DEPENDENCIES
@@ -667,13 +667,6 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
uv pip install --system ep_kernels/dist/*.whl --verbose \
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
# This must run after the vLLM wheel and EP kernels installs above, because
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
# (in the flashinfer-jit-cache layer) causes ~2.5 GB of layer duplication
# when a later pip install overwrites flashinfer package files.
RUN flashinfer show-config && flashinfer download-cubin
# CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
# return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
# consistently from the host (see https://github.com/vllm-project/vllm/issues/18859).
@@ -763,10 +756,6 @@ FROM vllm-base AS vllm-openai-base
ARG TARGETPLATFORM
ARG INSTALL_KV_CONNECTORS=false
ARG CUDA_VERSION
ARG VLLM_BUILD_COMMIT
ARG VLLM_BUILD_PIPELINE
ARG VLLM_BUILD_URL
ARG VLLM_IMAGE_TAG
ARG PIP_INDEX_URL UV_INDEX_URL
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -803,18 +792,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
fi
ENV VLLM_USAGE_SOURCE production-docker-image
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
VLLM_BUILD_URL=${VLLM_BUILD_URL:-} \
VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG:-local/vllm-openai:dev}
LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" \
org.opencontainers.image.revision="${VLLM_BUILD_COMMIT}" \
org.opencontainers.image.version="${VLLM_IMAGE_TAG}" \
org.opencontainers.image.url="${VLLM_BUILD_URL}" \
ai.vllm.build.commit="${VLLM_BUILD_COMMIT}" \
ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}" \
ai.vllm.build.url="${VLLM_BUILD_URL}" \
ai.vllm.image.tag="${VLLM_IMAGE_TAG}"
# define sagemaker first, so it is not default from `docker build`
FROM vllm-openai-base AS vllm-sagemaker
-1
View File
@@ -192,7 +192,6 @@ ADD ./tests/ ./tests/
ADD ./examples/ ./examples/
ADD ./benchmarks/ ./benchmarks/
ADD ./vllm/collect_env.py .
ADD ./docker/ ./docker/
ADD ./.buildkite/ ./.buildkite/
# install development dependencies (for testing)
+8 -15
View File
@@ -200,9 +200,9 @@ RUN cd /opt/rixl && \
# DeepEP build stage
FROM base AS build_deep
ARG ROCSHMEM_BRANCH="f0acb0c6"
ARG ROCSHMEM_BRANCH="ba0bf0f3"
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
ARG DEEPEP_BRANCH="a9ea9774"
ARG DEEPEP_BRANCH="5d90af8b"
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
ARG DEEPEP_NIC="cx7"
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
@@ -213,15 +213,18 @@ RUN git clone ${ROCSHMEM_REPO} \
&& git checkout ${ROCSHMEM_BRANCH} \
&& mkdir -p projects/rocshmem/build \
&& cd projects/rocshmem/build \
&& INSTALL_PREFIX=${ROCSHMEM_DIR} \
../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
&& bash ../scripts/build_configs/all_backends \
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-DROCM_PATH=/opt/rocm \
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-DUSE_EXTERNAL_MPI=OFF
# Build DeepEP wheel.
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
RUN git clone ${DEEPEP_REPO} \
&& cd DeepEP \
&& git checkout ${DEEPEP_BRANCH} \
&& python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
&& python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
# not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -385,16 +388,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
# above are not available once that RUN step completes.
COPY --from=export_vllm /*.whl /opt/vllm-wheels/
# Update rdma-core to support latest rocshmem
ARG DEEPEP_NIC
RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
cd /tmp/rdma-core && \
mkdir -p build && cd build && \
cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
+5 -14
View File
@@ -5,6 +5,9 @@ WORKDIR /workspace/
ARG PYTHON_VERSION=3.12
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
RUN apt clean && apt-get update -y && \
apt-get install -y --no-install-recommends --fix-missing \
curl \
@@ -23,20 +26,8 @@ RUN apt clean && apt-get update -y && \
python3.12-dev \
python3-pip
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
printf '%s\n' \
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
'Pin: version 2025.3*' \
'Pin-Priority: 1001' \
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
apt-get update -y && \
apt-get install -y --no-install-recommends \
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
intel-oneapi-mkl-devel-2025.3 \
intel-oneapi-dnnl-devel-2025.3 && \
rm -rf /var/lib/apt/lists/*
RUN apt update && apt upgrade -y && \
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
# Install UMD
RUN mkdir neo && \
+2 -28
View File
@@ -27,22 +27,6 @@ variable "COMMIT" {
default = ""
}
variable "VLLM_BUILD_COMMIT" {
default = "unknown"
}
variable "VLLM_BUILD_PIPELINE" {
default = "local"
}
variable "VLLM_BUILD_URL" {
default = ""
}
variable "VLLM_IMAGE_TAG" {
default = "local/vllm-openai:dev"
}
# Groups
group "default" {
@@ -62,10 +46,6 @@ target "_common" {
max_jobs = MAX_JOBS
nvcc_threads = NVCC_THREADS
torch_cuda_arch_list = TORCH_CUDA_ARCH_LIST
VLLM_BUILD_COMMIT = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
VLLM_BUILD_PIPELINE = VLLM_BUILD_PIPELINE
VLLM_BUILD_URL = VLLM_BUILD_URL
VLLM_IMAGE_TAG = VLLM_IMAGE_TAG
}
}
@@ -76,16 +56,10 @@ target "_labels" {
"org.opencontainers.image.title" = "vLLM"
"org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs"
"org.opencontainers.image.licenses" = "Apache-2.0"
"org.opencontainers.image.revision" = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
"org.opencontainers.image.version" = VLLM_IMAGE_TAG
"org.opencontainers.image.url" = VLLM_BUILD_URL
"ai.vllm.build.commit" = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
"ai.vllm.build.pipeline" = VLLM_BUILD_PIPELINE
"ai.vllm.build.url" = VLLM_BUILD_URL
"ai.vllm.image.tag" = VLLM_IMAGE_TAG
"org.opencontainers.image.revision" = COMMIT
}
annotations = [
"index,manifest:org.opencontainers.image.revision=${VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")}",
"index,manifest:org.opencontainers.image.revision=${COMMIT}",
]
}
Binary file not shown.

Before

Width:  |  Height:  |  Size: 156 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 182 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 193 KiB

+2 -2
View File
@@ -163,7 +163,7 @@ Running with a local file:
```bash
vllm run-batch \
-i features/openai_batch/openai_example_batch.jsonl \
-i offline_inference/openai_batch/openai_example_batch.jsonl \
-o results.jsonl \
--model meta-llama/Meta-Llama-3-8B-Instruct
```
@@ -172,7 +172,7 @@ Using remote file:
```bash
vllm run-batch \
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl \
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl \
-o results.jsonl \
--model meta-llama/Meta-Llama-3-8B-Instruct
```
+1 -1
View File
@@ -23,7 +23,7 @@ llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)
!!! note
With tensor parallelism enabled, each process will read the whole model and split it into chunks, which makes the disk reading time even longer (proportional to the size of tensor parallelism).
You can convert the model checkpoint to a sharded checkpoint using [examples/features/sharded_state/load_sharded_state_offline.py](../../examples/features/sharded_state/load_sharded_state_offline.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
You can convert the model checkpoint to a sharded checkpoint using [examples/offline_inference/save_sharded_state.py](../../examples/offline_inference/save_sharded_state.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
## Quantization
+1 -1
View File
@@ -42,7 +42,7 @@ Traces can be visualized using <https://ui.perfetto.dev/>.
#### Offline Inference
Refer to [examples/features/profiling/simple_profiling_offline.py](../../examples/features/profiling/simple_profiling_offline.py) for an example.
Refer to [examples/offline_inference/simple_profiling.py](../../examples/offline_inference/simple_profiling.py) for an example.
#### OpenAI Server
+1 -1
View File
@@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
+1 -1
View File
@@ -11,7 +11,7 @@ Automatic Prefix Caching (APC in short) caches the KV cache of existing queries,
Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an example:
[examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py](../../examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py)
[examples/offline_inference/automatic_prefix_caching.py](../../examples/offline_inference/automatic_prefix_caching.py)
## Example workloads
+2 -2
View File
@@ -6,12 +6,12 @@ This directory contains examples for extending the context length of models usin
## Offline Inference Example
The [`context_extension.py`](../../examples/features/context_extension/context_extension_offline.py) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat 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/features/context_extension/context_extension_offline.py
python examples/offline_inference/context_extension.py
```
## OpenAI Online Method
-54
View File
@@ -1,54 +0,0 @@
# IndexCache
IndexCache reduces redundant top-k computation in DeepSeek-V3.2 (DSA) models by caching and reusing top-k indices across layers.
## Background
DeepSeek-V3.2 uses a DeepSeek Sparse Attention (DSA) mechanism where top-k token selection is computed per layer. For deep models with many layers, this computation can be expensive. IndexCache allows skipping redundant top-k computations by reusing indices from previous layers.
See: [IndexCache Paper](https://arxiv.org/abs/2603.12201)
## Usage
### CLI
```bash
vllm serve deepseek-ai/DeepSeek-V3.2 \
--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
```
### Configuration Reference
| Parameter | Type | Default | Description |
|----------------------|------|---------|--------------------------------------------------------------------------------------------------------------------------------------------------|
| `use_index_cache` | bool | false | Enable IndexCache. Must be set to true to use this feature |
| `index_topk_freq` | int | 1 | Frequency (in layers) at which top-k is computed. 1 = compute on every layer (disabled), 4 = compute on 1/4 of layers |
| `index_topk_pattern` | str | null | Per-layer F/S pattern. Overrides index_topk_freq if set. Each character maps to one DSA layer: F = Full, S = Shared |
### Configuration Examples
**Using `index_topk_freq`** (compute every N layers):
```bash
vllm serve deepseek-ai/DeepSeek-V3.2 \
--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
```
**Using `index_topk_pattern`** (explicit per-layer control):
```bash
# custom pattern for 61 layers: F = compute, S = reuse
vllm serve deepseek-ai/DeepSeek-V3.2 \
--hf-overrides '{"use_index_cache": true, "index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSF"}'
```
## How It Works
1. When IndexCache is enabled, layers marked with `"F"` (Full) calculate and store top-k indices
2. Subsequent layers marked with `"S"` (Shared) receive the cached indices from the previous layer instead of recomputing
3. The cached indices are passed through the layer stack, reducing total computation
## Requirements
- DeepSeek-V3.2 or compatible DSA model
- `use_index_cache: true` via `--hf-overrides`
+1 -1
View File
@@ -47,7 +47,7 @@ the third parameter is the path to the LoRA adapter.
)
```
Check out [examples/features/lora/multilora_offline.py](../../examples/features/lora/multilora_offline.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
Check out [examples/offline_inference/multilora_inference.py](../../examples/offline_inference/multilora_inference.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
## Serving LoRA Adapters
+7 -7
View File
@@ -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.
@@ -958,7 +958,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.
+2 -2
View File
@@ -16,7 +16,7 @@ To input multi-modal data, follow this schema in [vllm.inputs.EmbedsPrompt][]:
You can pass prompt embeddings from Hugging Face Transformers models to the `'prompt_embeds'` field of the prompt embedding dictionary, as shown in the following examples:
[examples/features/prompt_embed/prompt_embed_offline.py](../../examples/features/prompt_embed/prompt_embed_offline.py)
[examples/offline_inference/prompt_embed_inference.py](../../examples/offline_inference/prompt_embed_inference.py)
## Online Serving
@@ -41,4 +41,4 @@ vllm serve meta-llama/Llama-3.2-1B-Instruct --runner generate \
Then, you can use the OpenAI client as follows:
[examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py](../../examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py)
[examples/online_serving/prompt_embed_inference_with_openai_client.py](../../examples/online_serving/prompt_embed_inference_with_openai_client.py)
-1
View File
@@ -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

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