diff --git a/.buildkite/ci_config.yaml b/.buildkite/ci_config.yaml index b199e554a739..a60a4194e9b3 100644 --- a/.buildkite/ci_config.yaml +++ b/.buildkite/ci_config.yaml @@ -8,8 +8,8 @@ run_all_patterns: - "CMakeLists.txt" - "requirements/common.txt" - "requirements/cuda.txt" - - "requirements/build.txt" - - "requirements/test.txt" + - "requirements/build/cuda.txt" + - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" diff --git a/.buildkite/ci_config_intel.yaml b/.buildkite/ci_config_intel.yaml index 375be84a396a..a1c0091e0f10 100644 --- a/.buildkite/ci_config_intel.yaml +++ b/.buildkite/ci_config_intel.yaml @@ -6,8 +6,8 @@ run_all_patterns: - "CMakeLists.txt" - "requirements/common.txt" - "requirements/xpu.txt" - - "requirements/build.txt" - - "requirements/test.txt" + - "requirements/build/cuda.txt" + - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index e466e2a52444..9b1044443780 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -46,7 +46,7 @@ steps: - tests/models/language/pooling/ commands: - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m " pytest -x -v -s tests/models/language/generation -m cpu_model pytest -x -v -s tests/models/language/pooling -m cpu_model" @@ -99,7 +99,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" - parallelism: 2 + parallelism: 3 - label: "Arm CPU Test" depends_on: [] diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh index 9131dfc71a0a..00ae34bba6d7 100755 --- a/.buildkite/image_build/image_build.sh +++ b/.buildkite/image_build/image_build.sh @@ -92,8 +92,8 @@ check_and_skip_if_image_exists() { } ecr_login() { - aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" - aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com + aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true + aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true } prepare_cache_tags() { diff --git a/.buildkite/image_build/image_build_cpu.sh b/.buildkite/image_build/image_build_cpu.sh index ccfe155fa2b7..035f070ab891 100755 --- a/.buildkite/image_build/image_build_cpu.sh +++ b/.buildkite/image_build/image_build_cpu.sh @@ -11,7 +11,7 @@ REPO=$2 BUILDKITE_COMMIT=$3 # authenticate with AWS ECR -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true # skip build if image already exists if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then diff --git a/.buildkite/image_build/image_build_cpu_arm64.sh b/.buildkite/image_build/image_build_cpu_arm64.sh index ff3d11c8d599..b561e2c2e463 100755 --- a/.buildkite/image_build/image_build_cpu_arm64.sh +++ b/.buildkite/image_build/image_build_cpu_arm64.sh @@ -11,7 +11,7 @@ REPO=$2 BUILDKITE_COMMIT=$3 # authenticate with AWS ECR -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true # skip build if image already exists if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh new file mode 100755 index 000000000000..cbd08aa7bd0b --- /dev/null +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -0,0 +1,68 @@ +#!/bin/bash +set -euo pipefail + +# Build a vLLM test image with PyTorch nightly installed. +# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group. + +if [[ $# -lt 5 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 +BRANCH=$4 +IMAGE_TAG=$5 + +# --- Arguments --- +echo "--- :mag: Arguments" +echo "REGISTRY: ${REGISTRY}" +echo "REPO: ${REPO}" +echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}" +echo "BRANCH: ${BRANCH}" +echo "IMAGE_TAG: ${IMAGE_TAG}" + +# --- ECR login --- +echo "--- :key: ECR login" +aws ecr-public get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com + +# --- Set up buildx --- +echo "--- :docker: Setting up buildx" +docker buildx create --name vllm-builder --driver docker-container --use || true +docker buildx inspect --bootstrap +docker buildx ls + +# --- Skip if image already exists --- +echo "--- :mag: Checking if image already exists" +if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then + echo "Image found: $IMAGE_TAG — skipping build" + exit 0 +fi +echo "Image not found, proceeding with build..." + +# --- CUDA 13.0 for nightly builds --- +# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9 +NIGHTLY_CUDA_VERSION="13.0.2" +NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04" +NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04" + +echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" +docker buildx build --file docker/Dockerfile \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg USE_SCCACHE=1 \ + --build-arg PYTORCH_NIGHTLY=1 \ + --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \ + --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ + --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ + --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ + --tag "$IMAGE_TAG" \ + --push \ + --target test \ + --progress plain . + +echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG" diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml new file mode 100644 index 000000000000..c66576d40991 --- /dev/null +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -0,0 +1,21 @@ +group: Engine Intel +depends_on: + - image-build-xpu +steps: +- label: Engine (1 GPU) + timeout_in_minutes: 30 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/v1/engine/ + - tests/v1/engine/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py' diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml new file mode 100644 index 000000000000..66a8db25f02e --- /dev/null +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -0,0 +1,21 @@ +group: Kernels Intel +depends_on: + - image-build-xpu +steps: +- label: vLLM IR Tests + timeout_in_minutes: 30 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ir + - vllm/kernels + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s kernels/ir' diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml new file mode 100644 index 000000000000..729edd159cdf --- /dev/null +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -0,0 +1,130 @@ +group: LoRA Intel +depends_on: + - image-build-xpu +steps: +- label: LoRA Runtime + Utils + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s lora/test_layers.py && + pytest -v -s lora/test_lora_checkpoints.py && + (pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) && + pytest -v -s lora/test_lora_huggingface.py && + pytest -v -s lora/test_lora_manager.py && + pytest -v -s lora/test_lora_utils.py && + pytest -v -s lora/test_peft_helper.py && + pytest -v -s lora/test_resolver.py && + pytest -v -s lora/test_utils.py && + (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) && + (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)' + +- label: LoRA Fused/MoE Kernels + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s lora/test_fused_moe_lora_kernel.py && + pytest -v -s lora/test_moe_lora_align_sum.py' + +- label: LoRA Punica Kernels + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + set -o pipefail && + pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"' + +- label: LoRA Punica FP8/XPU Ops + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s lora/test_punica_ops_fp8.py && + pytest -v -s lora/test_punica_xpu_ops.py' + +- label: LoRA Models + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) && + pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" && + pytest -v -s lora/test_qwen35_densemodel_lora.py && + pytest -v -s lora/test_transformers_model.py' + +- label: LoRA Multimodal + timeout_in_minutes: 45 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s lora/test_default_mm_loras.py && + (pytest -v -s lora/test_qwen3_unembed.py || true) && + (pytest -v -s lora/test_qwenvl.py || true) && + pytest -v -s lora/test_whisper.py' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml new file mode 100644 index 000000000000..39634ccccacf --- /dev/null +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -0,0 +1,55 @@ +group: Miscellaneous Intel +depends_on: + - image-build-xpu +steps: +- label: V1 Core + KV + Metrics + timeout_in_minutes: 30 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'pip install -r requirements/kv_connectors.txt && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + cd tests && + pytest -v -s v1/executor' + +- label: V1 Sample + Logits + timeout_in_minutes: 30 + device: intel_gpu + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && + cd tests && + pytest -v -s v1/logits_processors && + pytest -v -s v1/test_oracle.py && + pytest -v -s v1/test_request.py && + pytest -v -s v1/test_outputs.py' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 295840440ad5..c59be699502f 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -35,6 +35,7 @@ steps: python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 && + python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 && python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel' @@ -61,4 +62,4 @@ steps: pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py' + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py' diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 45b2996f7ead..8fce15680173 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -1,3 +1,13 @@ +# CUDA architecture lists — following PyTorch RELEASE.md +# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md) +# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support +env: + CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" + # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13) + CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX" + CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0" + CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0" + steps: - input: "Provide Release version here" id: input-release-version @@ -14,12 +24,10 @@ steps: agents: queue: arm64_cpu_queue_release commands: - # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: - # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "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" @@ -29,9 +37,7 @@ steps: agents: queue: arm64_cpu_queue_release commands: - # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: - # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "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 manylinux_2_35" @@ -57,7 +63,7 @@ steps: agents: queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --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 manylinux_2_31" @@ -70,7 +76,7 @@ steps: agents: queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "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 manylinux_2_35" @@ -98,99 +104,105 @@ steps: commands: - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + - block: "Unblock to build release Docker images" + depends_on: ~ + key: block-build-release-images + if: build.env("NIGHTLY") != "1" + - group: "Build release Docker images" key: "build-release-images" + depends_on: block-build-release-images + allow_dependency_failure: true steps: - - label: "Build release image - x86_64 - CUDA 12.9" + - label: "Build release image - x86_64 - CUDA 13.0" depends_on: ~ id: build-release-image-x86 agents: queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." + - "DOCKER_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" - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - label: "Build release image - aarch64 - CUDA 12.9" + - label: "Build release image - aarch64 - CUDA 13.0" depends_on: ~ id: build-release-image-arm64 agents: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." + - "DOCKER_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 13.0" + - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ - id: build-release-image-x86-cuda-13-0 + id: build-release-image-x86-cuda-12-9 agents: queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" + - "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)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - label: "Build release image - aarch64 - CUDA 13.0" + - label: "Build release image - aarch64 - CUDA 12.9" depends_on: ~ - id: build-release-image-arm64-cuda-13-0 + id: build-release-image-arm64-cuda-12-9 agents: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" + - "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 12.9 - Ubuntu 24.04" + - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04" depends_on: ~ id: build-release-image-x86-ubuntu2404 agents: queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "DOCKER_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" - - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04" + - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04" depends_on: ~ id: build-release-image-arm64-ubuntu2404 agents: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "DOCKER_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 13.0 - Ubuntu 24.04" + - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04" depends_on: ~ - id: build-release-image-x86-cuda-13-0-ubuntu2404 + id: build-release-image-x86-cuda-12-9-ubuntu2404 agents: queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404" - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404" + - "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" - - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04" + - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04" depends_on: ~ - id: build-release-image-arm64-cuda-13-0-ubuntu2404 + id: build-release-image-arm64-cuda-12-9-ubuntu2404 agents: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404" + - "DOCKER_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" key: block-cpu-release-image-build @@ -231,7 +243,7 @@ steps: - group: "Publish release images" key: "publish-release-images" steps: - - label: "Create multi-arch manifest - CUDA 12.9" + - label: "Create multi-arch manifest - CUDA 13.0" depends_on: - build-release-image-x86 - build-release-image-arm64 @@ -243,7 +255,7 @@ steps: - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend" - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - label: "Annotate release workflow - CUDA 12.9" + - label: "Annotate release workflow - CUDA 13.0" depends_on: - create-multi-arch-manifest id: annotate-release-workflow @@ -252,19 +264,19 @@ steps: commands: - "bash .buildkite/scripts/annotate-release.sh" - - label: "Create multi-arch manifest - CUDA 13.0" + - label: "Create multi-arch manifest - CUDA 12.9" depends_on: - - build-release-image-x86-cuda-13-0 - - build-release-image-arm64-cuda-13-0 - id: create-multi-arch-manifest-cuda-13-0 + - build-release-image-x86-cuda-12-9 + - build-release-image-arm64-cuda-12-9 + id: create-multi-arch-manifest-cuda-12-9 agents: queue: small_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04" + - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" depends_on: - build-release-image-x86-ubuntu2404 - build-release-image-arm64-ubuntu2404 @@ -276,17 +288,17 @@ steps: - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend" - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" + - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04" depends_on: - - build-release-image-x86-cuda-13-0-ubuntu2404 - - build-release-image-arm64-cuda-13-0-ubuntu2404 - id: create-multi-arch-manifest-cuda-13-0-ubuntu2404 + - build-release-image-x86-cuda-12-9-ubuntu2404 + - build-release-image-arm64-cuda-12-9-ubuntu2404 + id: create-multi-arch-manifest-cuda-12-9-ubuntu2404 agents: queue: small_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - label: "Publish nightly multi-arch image to DockerHub" depends_on: @@ -306,16 +318,16 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" - - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0" + - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9" depends_on: - - create-multi-arch-manifest-cuda-13-0 + - create-multi-arch-manifest-cuda-12-9 if: build.env("NIGHTLY") == "1" agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/push-nightly-builds.sh cu130" + - "bash .buildkite/scripts/push-nightly-builds.sh cu129" # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-" plugins: - docker-login#v3.0.0: username: vllmbot @@ -617,6 +629,8 @@ steps: - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: + - step: block-build-release-images + allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: diff --git a/.buildkite/scripts/annotate-release.sh b/.buildkite/scripts/annotate-release.sh index 2da9db2f2e5d..6f41d1cdda47 100755 --- a/.buildkite/scripts/annotate-release.sh +++ b/.buildkite/scripts/annotate-release.sh @@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF To download the wheel (by commit): \`\`\` -aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl . -aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl . +aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl . +aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl . -(Optional) For CUDA 13.0: -aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl . -aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl . +(Optional) For CUDA 12.9: +aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl . +aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl . (Optional) For CPU: aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl . @@ -33,8 +33,8 @@ To download and upload the image: docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 +docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 +docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} @@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 docker push vllm/vllm-openai:latest-x86_64 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130 -docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130 -docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 -docker push vllm/vllm-openai:latest-x86_64-cu130 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129 +docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 +docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 +docker push vllm/vllm-openai:latest-x86_64-cu129 +docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64 @@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 docker push vllm/vllm-openai:latest-aarch64 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130 -docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130 -docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 -docker push vllm/vllm-openai:latest-aarch64-cu130 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129 +docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 +docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 +docker push vllm/vllm-openai:latest-aarch64-cu129 +docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 ## ROCm @@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${ docker manifest push vllm/vllm-openai:latest docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} -docker manifest rm vllm/vllm-openai:latest-cu130 -docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130 -docker manifest push vllm/vllm-openai:latest-cu130 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130 +docker manifest rm vllm/vllm-openai:latest-cu129 +docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 +docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 +docker manifest push vllm/vllm-openai:latest-cu129 +docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 docker manifest rm vllm/vllm-openai-cpu:latest || true docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 diff --git a/.buildkite/scripts/check-ray-compatibility.sh b/.buildkite/scripts/check-ray-compatibility.sh index 1572fe94168d..b056d4403dbb 100644 --- a/.buildkite/scripts/check-ray-compatibility.sh +++ b/.buildkite/scripts/check-ray-compatibility.sh @@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then TORCH_INDEX_URL="" fi else - TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129" + TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130" fi echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}" diff --git a/.buildkite/scripts/generate-and-upload-nightly-index.sh b/.buildkite/scripts/generate-and-upload-nightly-index.sh index fa6eb979af53..88c4f5173139 100755 --- a/.buildkite/scripts/generate-and-upload-nightly-index.sh +++ b/.buildkite/scripts/generate-and-upload-nightly-index.sh @@ -9,7 +9,7 @@ set -ex BUCKET="vllm-wheels" INDICES_OUTPUT_DIR="indices" -DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py +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/" @@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) if [[ "$has_new_python" -eq 0 ]]; then # use new python from docker docker pull python:3-slim - PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3" + PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3" fi echo "Using python interpreter: $PYTHON" diff --git a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh index 4ff15067aa0f..f12bb524d4cb 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh @@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then exit 1 fi -#echo "--- DP+TP" -#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & -#server_pid=$! -#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 -#vllm bench serve \ -# --backend vllm \ -# --dataset-name random \ -# --model meta-llama/Llama-3.2-3B-Instruct \ -# --num-prompts 20 \ -# --result-dir ./test_results \ -# --result-filename dp_pp.json \ -# --save-result \ -# --endpoint /v1/completions -#kill -s SIGTERM $server_pid; wait $server_pid || true -#failed_req=$(jq '.failed' ./test_results/dp_pp.json) -#if [ "$failed_req" -ne 0 ]; then -# echo "Some requests were failed!" -# exit 1 -#fi +echo "--- DP+TP" +vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & +server_pid=$! +timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 +vllm bench serve \ + --backend vllm \ + --dataset-name random \ + --model meta-llama/Llama-3.2-3B-Instruct \ + --num-prompts 20 \ + --result-dir ./test_results \ + --result-filename dp_pp.json \ + --save-result \ + --endpoint /v1/completions +kill -s SIGTERM $server_pid; wait $server_pid || true +failed_req=$(jq '.failed' ./test_results/dp_pp.json) +if [ "$failed_req" -ne 0 ]; then + echo "Some requests were failed!" + exit 1 +fi diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index c2509a07b2c4..7166435ac1e9 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -51,6 +51,7 @@ function cpu_tests() { set -e pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/attention/test_cpu_attn.py + pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic" # basic online serving diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index db75ad3083b2..27ec0068668f 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image" docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu . # Run the image, setting --shm-size=4g for tensor parallel. -docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \ +docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \ timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" diff --git a/.buildkite/scripts/hardware_ci/run-hpu-test.sh b/.buildkite/scripts/hardware_ci/run-hpu-test.sh index 10df07b2000f..0b5d0af4b6f3 100644 --- a/.buildkite/scripts/hardware_ci/run-hpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-hpu-test.sh @@ -42,7 +42,7 @@ WORKDIR /workspace/vllm ENV no_proxy=localhost,127.0.0.1 ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true -RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build.txt)' +RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build/cuda.txt)' RUN VLLM_TARGET_DEVICE=empty pip install --no-build-isolation -e . RUN pip install git+https://github.com/vllm-project/vllm-gaudi.git diff --git a/.buildkite/scripts/hardware_ci/run-intel-test.sh b/.buildkite/scripts/hardware_ci/run-intel-test.sh index 836f0b849573..eae3b231b4be 100755 --- a/.buildkite/scripts/hardware_ci/run-intel-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-test.sh @@ -25,22 +25,100 @@ export PYTHONPATH=".." ############################################################################### cleanup_docker() { + # Share the same lock with image pull to avoid cleanup/pull races on one node. + local docker_lock="/tmp/docker-pull.lock" + exec 9>"$docker_lock" + flock 9 + docker_root=$(docker info -f '{{.DockerRootDir}}') if [ -z "$docker_root" ]; then echo "Failed to determine Docker root directory." >&2 - exit 1 + flock -u 9 + return 1 fi echo "Docker root directory: $docker_root" disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') threshold=70 if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - docker image prune -f - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." + echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..." + cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1 + else + echo "Disk usage is below $threshold%. Checking old CI images anyway." + cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0 + fi + echo "Old CI image cleanup completed." + + flock -u 9 +} + +cleanup_old_ci_images() { + local repo_prefix="$1" + local current_image_ref="$2" + local ttl_hours="$3" + local aggressive_cleanup="$4" + + if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then + echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'" + return 0 + fi + + if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then + echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72" + ttl_hours=72 + fi + + local now_epoch cutoff_epoch + now_epoch=$(date +%s) + cutoff_epoch=$((now_epoch - ttl_hours * 3600)) + + local -a used_image_ids + mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u) + + local removed_count=0 + local examined_count=0 + declare -A seen_ids=() + + while read -r image_ref image_id; do + [[ -z "$image_ref" || -z "$image_id" ]] && continue + ((examined_count++)) + + # Keep the image this job is going to use. + if [[ "$image_ref" == "$current_image_ref" ]]; then + continue + fi + + # Avoid duplicate deletes when multiple tags point to same image id. + if [[ -n "${seen_ids[$image_id]:-}" ]]; then + continue + fi + seen_ids[$image_id]=1 + + # Never delete images that are used by any container on this node. + if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then + continue + fi + + local created created_epoch + created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true) + [[ -z "$created" ]] && continue + created_epoch=$(date -d "$created" +%s 2>/dev/null || true) + [[ -z "$created_epoch" ]] && continue + + if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then + if docker image rm -f "$image_id" >/dev/null 2>&1; then + ((removed_count++)) + fi + fi + done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}') + + # Also trim old dangling layers; this is safe and does not remove referenced images. + docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true + + if [[ "$aggressive_cleanup" == "1" ]]; then + echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure." else - echo "Disk usage is below $threshold%. No cleanup needed." + echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)." fi } @@ -240,7 +318,6 @@ fi cleanup_docker aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" -aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com # --- Build or pull test image --- IMAGE="${IMAGE_TAG_XPU:-${image_name}}" @@ -266,8 +343,6 @@ fi remove_docker_container() { docker rm -f "${container_name}" || true - docker image rm -f "${image_name}" || true - docker system prune -f || true } trap remove_docker_container EXIT @@ -283,6 +358,7 @@ docker run \ --ipc=host \ --privileged \ -v /dev/dri/by-path:/dev/dri/by-path \ + -v ${HOME}/.cache/huggingface:/root/.cache/huggingface \ --entrypoint="" \ -e "HF_TOKEN=${HF_TOKEN:-}" \ -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \ diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index 0467da7c2307..14bd08cfc1c4 100644 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -12,9 +12,7 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu . # Setup cleanup remove_docker_container() { - docker rm -f "${container_name}" || true; - docker image rm -f "${image_name}" || true; - docker system prune -f || true; + docker rm -f "${container_name}" || true } trap remove_docker_container EXIT @@ -50,6 +48,6 @@ docker run \ pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py pytest -v -s v1/structured_output pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)" + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py pytest -v -s v1/test_serial_utils.py ' diff --git a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh new file mode 100755 index 000000000000..06743f16b687 --- /dev/null +++ b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euxo pipefail + +# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] +THRESHOLD=${1:-0.8} +NUM_Q=${2:-1319} +PORT=${3:-8050} +OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled} +mkdir -p "${OUT_DIR}" + +wait_for_server() { + local port=$1 + timeout 600 bash -c ' + until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do + sleep 1 + done' +} + +MODEL="Qwen/Qwen3-30B-A3B-FP8" +BACK="allgather_reducescatter" + +cleanup() { + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "${SERVER_PID}" 2>/dev/null || break + sleep 0.5 + done + kill -9 "${SERVER_PID}" 2>/dev/null || true + fi +} +trap cleanup EXIT + +VLLM_DEEP_GEMM_WARMUP=skip \ +vllm serve "$MODEL" \ +--enforce-eager \ +--data-parallel-size 4 \ +--enable-expert-parallel \ +--enable-eplb \ +--all2all-backend "$BACK" \ +--eplb-config '{"window_size":20, "step_interval":100, "use_async":true}' \ +--trust-remote-code \ +--max-model-len 2048 \ +--port "$PORT" & +SERVER_PID=$! +wait_for_server "$PORT" + +TAG=$(echo "$MODEL" | tr '/: \\n' '_____') +OUT="${OUT_DIR}/${TAG}_${BACK}.json" +python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}" +python3 - <= ${THRESHOLD}, f"${MODEL} ${BACK} accuracy {acc}" +PY diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c7e9e13f2fd3..a8c9e4094387 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1,8 +1,8 @@ # In this file, you can add more tests to run either by adding a new step or # adding a new command to an existing step. See different options here for examples. -# This script will be feed into Jinja template in `test-template-aws.j2` at -# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-aws.j2 +# This script will be feed into Jinja template in `test-template-amd.j2` at +# https://github.com/vllm-project/buildkite-ci/blob/main/scripts/test-template-amd.j2 # to generate the final pipeline yaml file. # Documentation @@ -39,8 +39,8 @@ ##################################################################################################################################### # # # IMPORTANT: # -# * Currently AMD CI has MI250 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are tracked in: # -# https://github.com/vllm-project/vllm/issues/34994 # +# * Currently AMD CI has MI250 agents, MI300 agents, MI325 agents, and MI355 agents. All upcoming feature improvements are # +# tracked in: https://github.com/vllm-project/vllm/issues/34994 # # # #-----------------------------------------------------------------------------------------------------------------------------------# # # @@ -104,193 +104,195 @@ ##################################################################################################################################### - - steps: +######################################################################################################################################### +# # +# MI250 (gfx90a) tests # +# # +######################################################################################################################################### -##################################################################################################################################### -# # -# MI250 test definitions ( currently the test set is completely mirrored // TBD which tests are to be routed there ultimately) # -# # -##################################################################################################################################### +#----------------------------------------------------- mi250 · basic_correctness -----------------------------------------------------# -- label: Pytorch Nightly Dependency Override Check # TBD +- label: Distributed Model Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - soft_fail: true + agent_pool: mi250_2 + num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/nightly_torch_test.txt + - vllm/model_executor/model_loader/sharded_state_loader.py + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/basic_correctness/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + - tests/models/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - bash standalone_tests/pytorch_nightly_dependency.sh + - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' +#-------------------------------------------------------- mi250 · benchmarks ---------------------------------------------------------# -- label: Async Engine, Inputs, Utils, Worker # TBD +- label: Benchmarks # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/.buildkite" source_file_dependencies: - - vllm/ - - tests/detokenizer - - tests/multimodal - - tests/utils_ + - benchmarks/ + - vllm/platforms/rocm.py commands: - - pytest -v -s detokenizer - - pytest -v -s -m 'not cpu_test' multimodal - - pytest -v -s utils_ + - bash scripts/run-benchmarks.sh +#---------------------------------------------------------- mi250 · compile ----------------------------------------------------------# -- label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD +- label: PyTorch Compilation Unit Tests # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + torch_nightly: true optional: true - no_gpu: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_inputs.py - - tests/test_outputs.py - - tests/test_pooling_params.py - - tests/test_ray_env.py - - tests/multimodal - - tests/renderers - - tests/standalone_tests/lazy_imports.py - - tests/tokenizers_ - - tests/tool_parsers - - tests/transformers_utils - - tests/config + - vllm/compilation/ + - vllm/model_executor/layers/ + - vllm/v1/worker/ + - vllm/v1/attention/ + - vllm/v1/cudagraph_dispatcher.py + - vllm/config/compilation.py + - csrc/ + - tests/compile + - vllm/platforms/rocm.py commands: - - python3 standalone_tests/lazy_imports.py - - pytest -v -s test_inputs.py - - pytest -v -s test_outputs.py - - pytest -v -s test_pooling_params.py - - pytest -v -s test_ray_env.py - - pytest -v -s -m 'cpu_test' multimodal - - pytest -v -s renderers - - pytest -v -s tokenizers_ - - pytest -v -s tool_parsers - - pytest -v -s transformers_utils - - pytest -v -s config - + - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" -- label: Python-only Installation # TBD +- label: PyTorch Fullgraph # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - tests/standalone_tests/python_only_compile.sh - - setup.py + - vllm/compilation/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/config/compilation.py + - csrc/ + - tests/compile - vllm/platforms/rocm.py commands: - - bash standalone_tests/python_only_compile.sh - + - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' -- label: Basic Correctness # TBD +- label: PyTorch Fullgraph Smoke Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - fast_check: true + optional: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/basic_correctness/test_basic_correctness - - tests/basic_correctness/test_cpu_offload - - tests/basic_correctness/test_cumem.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s basic_correctness/test_cumem.py - - pytest -v -s basic_correctness/test_basic_correctness.py - - pytest -v -s basic_correctness/test_cpu_offload.py - - -- label: Entrypoints Unit Tests # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/entrypoints - - tests/entrypoints/ + - vllm/compilation/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/config/compilation.py + - csrc/ + - tests/compile - vllm/platforms/rocm.py commands: - - pytest -v -s entrypoints/openai/tool_parsers - - 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 - + - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;" -- label: Entrypoints Integration (LLM) # TBD +- label: Distributed Compile + RPC Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true + agent_pool: mi250_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/entrypoints/llm - - tests/entrypoints/offline_mode + - vllm/compilation/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/compile/fullgraph/test_basic_correctness.py + - tests/compile/test_wrapper.py + - tests/entrypoints/llm/test_collective_rpc.py + - vllm/platforms/rocm.py commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py - - pytest -v -s entrypoints/llm/test_generate.py - - pytest -v -s entrypoints/offline_mode + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s entrypoints/llm/test_collective_rpc.py + - pytest -v -s ./compile/fullgraph/test_basic_correctness.py + - pytest -v -s ./compile/test_wrapper.py +#-------------------------------------------------------- mi250 · distributed --------------------------------------------------------# -- label: Entrypoints Integration (API Server 2) # TBD +- label: Distributed Comm Ops # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true + agent_pool: mi250_2 + num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - vllm/distributed + - tests/distributed + - vllm/platforms/rocm.py commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use - + - pytest -v -s distributed/test_comm_ops.py + - pytest -v -s distributed/test_shm_broadcast.py + - pytest -v -s distributed/test_shm_buffer.py + - pytest -v -s distributed/test_shm_storage.py -- label: Entrypoints Integration (Responses API) # TBD +- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - fast_check: true - torch_nightly: true + agent_pool: mi250_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/entrypoints/openai/responses + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + - vllm/platforms/rocm.py commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/responses - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown + - pytest -v -s v1/worker/test_worker_memory_snapshot.py -- label: EPLB Algorithm # TBD +- label: Elastic EP Scaling Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + agent_pool: mi250_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/eplb - - tests/distributed/test_eplb_algo.py + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/compilation/ + - tests/distributed/ - vllm/platforms/rocm.py commands: - - pytest -v -s distributed/test_eplb_algo.py - + - pytest -v -s distributed/test_elastic_ep.py - label: EPLB Execution # TBD timeout_in_minutes: 180 @@ -307,8 +309,7 @@ steps: - pytest -v -s distributed/test_eplb_execute.py - pytest -v -s distributed/test_eplb_spec_decode.py - -- label: Elastic EP Scaling Test # TBD +- label: Pipeline + Context Parallelism (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_4 @@ -318,1177 +319,1020 @@ steps: - vllm/distributed/ - vllm/engine/ - vllm/executor/ - - vllm/compilation/ + - vllm/model_executor/models/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py - tests/distributed/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s distributed/test_elastic_ep.py + - pytest -v -s distributed/test_pp_cudagraph.py + - pytest -v -s distributed/test_pipeline_parallel.py +#---------------------------------------------------------- mi250 · engine -----------------------------------------------------------# -- label: Metrics, Tracing (2 GPUs) # TBD +- label: Engine # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/v1/tracing + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port commands: - - "pip install \ - 'opentelemetry-sdk>=1.26.0' \ - 'opentelemetry-api>=1.26.0' \ - 'opentelemetry-exporter-otlp>=1.26.0' \ - 'opentelemetry-semantic-conventions-ai>=0.4.1'" - - pytest -v -s v1/tracing + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py +#----------------------------------------------------------- mi250 · evals -----------------------------------------------------------# -- label: Regression # TBD +- label: Multi-Modal Accuracy Eval (Small Models) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" + optional: true + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/ - - tests/test_regression + - vllm/multimodal/ + - vllm/inputs/ + - vllm/v1/core/ + - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ commands: - - pip install modelscope - - pytest -v -s test_regression.py + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 +#--------------------------------------------------------- mi250 · examples ----------------------------------------------------------# -- label: Engine # TBD +- label: Examples # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/examples" source_file_dependencies: - - vllm/ - - tests/engine - - tests/test_sequence - - tests/test_config - - tests/test_logger - - tests/test_vllm_port + - vllm/entrypoints + - vllm/multimodal + - examples/ + - vllm/platforms/rocm.py commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py + - pip install tensorizer + # Basic + - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN + - python3 basic/offline_inference/generate.py --model facebook/opt-125m + - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 + - python3 basic/offline_inference/classify.py + - python3 basic/offline_inference/embed.py + - python3 basic/offline_inference/score.py + # Multi-modal models + - 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 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 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 ----------------------------------------------------------# -- label: Engine (1 GPU) # TBD +- label: Kernels Core Operation Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/ - - tests/v1/engine/ + - csrc/ + - tests/kernels/core + - tests/kernels/test_top_k_per_row.py + - tests/kernels/test_concat_mla_q.py + - vllm/model_executor/layers/rotary_embedding/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/engine/test_preprocess_error_handling.py - - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py - + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py -- label: e2e Scheduling (1 GPU) # TBD +- label: Kernels Helion Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ + - vllm/utils/import_utils.py + - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/general/test_async_scheduling.py - + - pip install helion==1.0.0 + - pytest -v -s kernels/helion/ -- label: e2e Core (1 GPU) # TBD +- label: Kernels Mamba Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ + - csrc/mamba/ + - tests/kernels/mamba + - vllm/model_executor/layers/mamba/ops - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + - pytest -v -s kernels/mamba +#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------# -- label: Spec Decode Speculators + MTP # TBD +- label: LoRA %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + parallelism: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ + - vllm/lora + - tests/lora - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py +#------------------------------------------------------ mi250 · model_executor -------------------------------------------------------# -- label: Spec Decode Ngram + Suffix # TBD +- label: Model Executor # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ + - vllm/engine/arg_utils.py + - vllm/config/model.py + - vllm/model_executor + - tests/model_executor + - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + - apt-get update && apt-get install -y curl libsodium23 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s model_executor -m '(not slow_test)' + - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py +#---------------------------------------------------------- mi250 · models -----------------------------------------------------------# -- label: Spec Decode Draft Model # TBD +- label: Basic Models Test (Other CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + no_gpu: true optional: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ + - vllm/ + - tests/models/test_utils.py + - tests/models/test_vision.py + commands: + - pytest -v -s models/test_utils.py models/test_vision.py + +- label: Basic Models Tests (Extra Initialization) %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_1 + torch_nightly: true + parallelism: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/models/ - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ + - tests/models/test_initialization.py + - tests/models/registry.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" - + - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: V1 e2e (2 GPUs) # TBD +- label: Basic Models Tests (Initialization) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/v1/e2e + - tests/models/test_initialization.py + - tests/models/registry.py commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" - + - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 60 +- label: Basic Models Tests (Other) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py + - tests/models/test_terratorch.py + - tests/models/test_transformers.py + - tests/models/test_registry.py commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - + - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 60 +- label: Language Models Test (MTEB) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py + - tests/models/language/pooling_mteb_test commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - + - pytest -v -s models/language/pooling_mteb_test -- label: V1 Speculative Decoding (slow) # TBD +- label: Language Models Test (PPL) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ - - vllm/platforms/rocm.py + - vllm/ + - tests/models/language/generation_ppl_test commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - + - pytest -v -s models/language/generation_ppl_test -- label: V1 attention (H100-MI250) # TBD +- label: Language Models Tests (Extra Standard) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + torch_nightly: true + parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/models/language/pooling/test_embedding.py + - tests/models/language/generation/test_common.py + - tests/models/language/pooling/test_classification.py - vllm/_aiter_ops.py - - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention - + - pip freeze | grep -E 'torch' + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: V1 others (CPU) # TBD +- label: Multi-Modal Models (Extended Generation 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - no_gpu: true - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/v1 + - tests/models/multimodal/generation commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics - + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' -- label: Examples # TBD +- label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/examples" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/entrypoints - - vllm/multimodal - - examples/ - - vllm/platforms/rocm.py + - vllm/ + - tests/models/multimodal/pooling commands: - - pip install tensorizer - # Basic - - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN - - python3 basic/offline_inference/generate.py --model facebook/opt-125m - - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 - - python3 basic/offline_inference/classify.py - - python3 basic/offline_inference/embed.py - - python3 basic/offline_inference/score.py - # Multi-modal models - - 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 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 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 - + - pytest -v -s models/multimodal/pooling -m 'not core_model' -- label: Platform Tests (CUDA) # TBD +- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/cuda + - tests/models/multimodal commands: - - pytest -v -s cuda/test_cuda_context.py - - pytest -v -s cuda/test_platform_no_cuda_init.py + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" + - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model +#---------------------------------------------------------- mi250 · plugins ----------------------------------------------------------# -- label: Samplers Test # TBD +- label: Plugin Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + agent_pool: mi250_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/layers - - vllm/sampling_metadata.py - - vllm/v1/sample/ - - vllm/beam_search.py - - tests/samplers - - tests/conftest.py - - vllm/_aiter_ops.py + - vllm/plugins/ + - tests/plugins/ - vllm/platforms/rocm.py commands: - - pytest -v -s samplers + # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform + - pip install -e ./plugins/vllm_add_dummy_platform + - pytest -v -s plugins_tests/test_platform_plugins.py + - pip uninstall vllm_add_dummy_platform -y + # END: platform plugin tests + # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin + - pip install -e ./plugins/prithvi_io_processor_plugin + - pytest -v -s plugins_tests/test_io_processor_plugins.py + - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py + - pip uninstall prithvi_io_processor_plugin -y + # END: `io_processor` plugins test + # BEGIN: `bge_m3_sparse io_processor` test + - pip install -e ./plugins/bge_m3_sparse_plugin + - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py + - pip uninstall bge_m3_sparse_plugin -y + # END: `bge_m3_sparse io_processor` test + # BEGIN: `stat_logger` plugins test + - pip install -e ./plugins/vllm_add_dummy_stat_logger + - pytest -v -s plugins_tests/test_stats_logger_plugins.py + - pip uninstall dummy_stat_logger -y + # END: `stat_logger` plugins test + # BEGIN: other tests + - pytest -v -s plugins_tests/test_scheduler_plugins.py + - pip install -e ./plugins/vllm_add_dummy_model + - pytest -v -s distributed/test_distributed_oot.py + - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py + - pytest -v -s models/test_oot_registration.py + - pytest -v -s plugins/lora_resolvers +#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# -- label: LoRA %N # TBD +- label: Batch Invariance (H100-MI250) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - parallelism: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/lora - - tests/lora + - vllm/v1/attention + - vllm/model_executor/layers + - tests/v1/determinism/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py - + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -- label: PyTorch Compilation Unit Tests # TBD +- label: Cudagraph # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - torch_nightly: true optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/ - - vllm/v1/worker/ - - vllm/v1/attention/ + - tests/v1/cudagraph - vllm/v1/cudagraph_dispatcher.py - vllm/config/compilation.py - - csrc/ - - tests/compile + - vllm/compilation - vllm/platforms/rocm.py commands: - - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" - + - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py + - pytest -v -s v1/cudagraph/test_cudagraph_mode.py -- label: PyTorch Fullgraph Smoke Test # TBD +- label: e2e Core (1 GPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/ - - vllm/v1/attention/ - - vllm/config/compilation.py - - csrc/ - - tests/compile + - vllm/v1/ + - tests/v1/e2e/general/ - vllm/platforms/rocm.py commands: - - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;" - + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py -- label: PyTorch Fullgraph # TBD +- label: e2e Scheduling (1 GPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/ - - vllm/v1/attention/ - - vllm/config/compilation.py - - csrc/ - - tests/compile + - vllm/v1/ + - tests/v1/e2e/general/ - vllm/platforms/rocm.py commands: - - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile' - + - pytest -v -s v1/e2e/general/test_async_scheduling.py -- label: Cudagraph # TBD +- label: Engine (1 GPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - tests/v1/cudagraph - - vllm/v1/cudagraph_dispatcher.py - - vllm/config/compilation.py - - vllm/compilation + - vllm/v1/ + - tests/v1/engine/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - - pytest -v -s v1/cudagraph/test_cudagraph_mode.py - + - pytest -v -s v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py -- label: Kernels Core Operation Test # TBD +- label: Spec Decode Draft Model # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - tests/kernels/core - - tests/kernels/test_top_k_per_row.py - - tests/kernels/test_concat_mla_q.py - - vllm/model_executor/layers/rotary_embedding/ - - vllm/_aiter_ops.py + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/core kernels/test_top_k_per_row.py - + - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" -- label: Kernels Mamba Test # TBD +- label: Spec Decode Speculators + MTP # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/mamba/ - - tests/kernels/mamba - - vllm/model_executor/layers/mamba/ops + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/mamba - + - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" -- label: Kernels Helion Test # TBD +- label: V1 attention (H100-MI250) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/utils/import_utils.py - - tests/kernels/helion/ + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention + - vllm/_aiter_ops.py + - vllm/envs.py - vllm/platforms/rocm.py commands: - - pip install helion==0.3.3 - - pytest -v -s kernels/helion/ - + - pytest -v -s v1/attention -- label: Model Executor # TBD - timeout_in_minutes: 180 +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 optional: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/engine/arg_utils.py - - vllm/config/model.py - - vllm/model_executor - - tests/model_executor - - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - apt-get update && apt-get install -y curl libsodium23 - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s model_executor -m '(not slow_test)' - - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py - - -- label: Benchmarks # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/.buildkite" - source_file_dependencies: - - benchmarks/ - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py commands: - - bash scripts/run-benchmarks.sh - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: Benchmarks CLI Test # TBD - timeout_in_minutes: 180 +- label: V1 Sample + Logits # TBD + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/benchmarks/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py commands: - - pytest -v -s benchmarks/ - + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py -- label: OpenAI API correctness # TBD +- label: Distributed DP Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" + agent_pool: mi250_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/entrypoints/openai/ - - vllm/model_executor/models/whisper.py - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/v1/distributed + - tests/entrypoints/openai/test_multi_api_servers.py - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ commands: - - bash ../tools/install_torchcodec_rocm.sh || exit 1 - - pytest -s entrypoints/openai/correctness/ + - export TORCH_NCCL_BLOCKING_WAIT=1 + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py +- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + agent_pool: mi250_2 + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh -- label: Basic Models Tests (Initialization) # TBD +- label: V1 e2e (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true + agent_pool: mi250_2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/test_initialization.py - - tests/models/registry.py + - tests/v1/e2e commands: - - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" -- label: Basic Models Tests (Extra Initialization) %N # TBD +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - parallelism: 2 + agent_pool: mi250_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - tests/models/test_initialization.py - - tests/models/registry.py - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# -- label: Basic Models Tests (Other) # TBD +- label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 - torch_nightly: true + optional: true + no_gpu: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/test_terratorch.py - - tests/models/test_transformers.py - - tests/models/test_registry.py + - tests/test_inputs.py + - tests/test_outputs.py + - tests/test_pooling_params.py + - tests/test_ray_env.py + - tests/multimodal + - tests/renderers + - tests/standalone_tests/lazy_imports.py + - tests/tokenizers_ + - tests/tool_parsers + - tests/transformers_utils + - tests/config commands: - - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + - python3 standalone_tests/lazy_imports.py + - pytest -v -s test_inputs.py + - pytest -v -s test_outputs.py + - pytest -v -s test_pooling_params.py + - pytest -v -s test_ray_env.py + - pytest -v -s -m 'cpu_test' multimodal + - pytest -v -s renderers + - pytest -v -s tokenizers_ + - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py + - pytest -v -s tool_parsers + - pytest -v -s transformers_utils + - pytest -v -s config +######################################################################################################################################### +# # +# MI300 (gfx942) tests # +# # +######################################################################################################################################### -- label: Basic Models Test (Other CPU) # TBD +#----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------# + +- label: Basic Correctness # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true - optional: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/test_utils.py - - tests/models/test_vision.py + - tests/basic_correctness/test_basic_correctness + - tests/basic_correctness/test_cpu_offload + - tests/basic_correctness/test_cumem.py commands: - - pytest -v -s models/test_utils.py models/test_vision.py - + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s basic_correctness/test_cumem.py + - pytest -v -s basic_correctness/test_basic_correctness.py + - pytest -v -s basic_correctness/test_cpu_offload.py -- label: Language Models Tests (Extra Standard) %N # TBD +- label: Distributed Model Tests (2 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - parallelism: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/model_executor/model_loader/sharded_state_loader.py - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - vllm/model_executor/layers/ - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - - tests/models/language/pooling/test_embedding.py - - tests/models/language/generation/test_common.py - - tests/models/language/pooling/test_classification.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/basic_correctness/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + - tests/models/ commands: - - pip freeze | grep -E 'torch' - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py + - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' +#-------------------------------------------------------- mi300 · benchmarks ---------------------------------------------------------# -- label: Language Models Test (PPL) # TBD +- label: Benchmarks # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/.buildkite" + source_file_dependencies: + - benchmarks/ + - vllm/platforms/rocm.py + commands: + - bash scripts/run-benchmarks.sh + +- label: Benchmarks CLI Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/language/generation_ppl_test + - tests/benchmarks/ commands: - - pytest -v -s models/language/generation_ppl_test + - pytest -v -s benchmarks/ +#---------------------------------------------------------- mi300 · compile ----------------------------------------------------------# -- label: Language Models Test (Extended Pooling) # TBD +- label: Fusion E2E Config Sweep (H100-MI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace/" source_file_dependencies: - - vllm/ - - tests/models/language/pooling + - csrc/quantization/ + - vllm/compilation/ + - vllm/model_executor/layers/layernorm.py + - vllm/model_executor/layers/activation.py + - vllm/model_executor/layers/attention/attention.py + - vllm/model_executor/layers/quantization/input_quant_fp8.py + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s models/language/pooling -m 'not core_model' + - rocm-smi + - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" +- label: Fusion E2E Quick (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/compilation/ + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - rocm-smi + # Run all models and attn backends but only Inductor partition and native custom ops + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'" + # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER + - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'" -- label: Language Models Test (MTEB) # TBD +- label: PyTorch Compilation Passes Unit Tests # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/language/pooling_mteb_test + - tests/compile/passes commands: - - pytest -v -s models/language/pooling_mteb_test - + - pytest -s -v compile/passes --ignore compile/passes/distributed -- label: Multi-Modal Processor (CPU) # TBD +- label: Pytorch Nightly Dependency Override Check # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - no_gpu: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true + soft_fail: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal - - tests/models/registry.py + - requirements/test/nightly-torch.txt + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py - + - bash standalone_tests/pytorch_nightly_dependency.sh -- label: Multi-Modal Accuracy Eval (Small Models) # TBD +- label: Distributed Compile Unit Tests (2xH100-2xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/" source_file_dependencies: - - vllm/multimodal/ - - vllm/inputs/ - - vllm/v1/core/ + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/passes/distributed/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ commands: - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py + - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py +#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------# -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD +- label: Platform Tests (CUDA) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal + - tests/cuda commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model + - pytest -v -s cuda/test_cuda_context.py + - pytest -v -s cuda/test_platform_no_cuda_init.py +#-------------------------------------------------------- mi300 · detokenizer --------------------------------------------------------# -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD +- label: Async Engine, Inputs, Utils, Worker # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - - -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model - - -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - torch_nightly: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - - -- label: Multi-Modal Models (Extended Generation 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py + - tests/detokenizer + - tests/multimodal + - tests/utils_ commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py + - pytest -v -s detokenizer + - pytest -v -s -m 'not cpu_test' multimodal + - pytest -v -s utils_ +#-------------------------------------------------------- mi300 · distributed --------------------------------------------------------# -- label: Multi-Modal Models (Extended Generation 2) # TBD +- label: EPLB Algorithm # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation + - vllm/distributed/eplb + - tests/distributed/test_eplb_algo.py + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' - + - pytest -v -s distributed/test_eplb_algo.py + - pytest -v -s distributed/test_eplb_utils.py -- label: Multi-Modal Models (Extended Generation 3) # TBD +- label: Distributed Tests (2xH100-2xMI250) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + working_dir: "/vllm-workspace/" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation + - vllm/distributed/ + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/distributed/test_context_parallel.py + - examples/offline_inference/data_parallel.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s tests/distributed/test_context_parallel.py + - 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: Multi-Modal Models (Extended Pooling) # TBD +- label: Distributed Tests (4xA100-4xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/pooling commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s distributed/test_custom_all_reduce.py + - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py + - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - pytest -v -s -x lora/test_mixtral.py -- label: Distributed Comm Ops # TBD +- label: Distributed Torchrun + Examples (4 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed - - tests/distributed + - vllm/distributed/ + - tests/distributed/test_torchrun_example.py + - tests/distributed/test_torchrun_example_moe.py + - examples/rl/ + - tests/examples/offline_inference/data_parallel.py - vllm/platforms/rocm.py commands: - - pytest -v -s distributed/test_comm_ops.py - - pytest -v -s distributed/test_shm_broadcast.py - - pytest -v -s distributed/test_shm_buffer.py - - pytest -v -s distributed/test_shm_storage.py - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py + - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py + - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py + - 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/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 (2 GPUs) # TBD +- label: Elastic EP Scaling Test # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/ - vllm/engine/ - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/v1/distributed - - tests/entrypoints/openai/test_multi_api_servers.py + - vllm/compilation/ + - tests/distributed/ - vllm/platforms/rocm.py commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py - + - pytest -v -s distributed/test_elastic_ep.py -- label: Distributed Compile + RPC Tests (2 GPUs) # TBD +- label: RayExecutorV2 (4 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/compilation/ - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/compile/fullgraph/test_basic_correctness.py - - tests/compile/test_wrapper.py - - tests/entrypoints/llm/test_collective_rpc.py + - vllm/v1/executor/ray_executor_v2.py + - vllm/v1/executor/abstract.py + - vllm/v1/executor/multiproc_executor.py + - tests/distributed/test_ray_v2_executor.py + - tests/distributed/test_ray_v2_executor_e2e.py + - tests/distributed/test_pipeline_parallel.py + - tests/basic_correctness/test_basic_correctness.py - vllm/platforms/rocm.py commands: + - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s entrypoints/llm/test_collective_rpc.py - - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - - pytest -v -s ./compile/test_wrapper.py + - pytest -v -s distributed/test_ray_v2_executor.py + - pytest -v -s distributed/test_ray_v2_executor_e2e.py + - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" + - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray" - -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD +- label: Distributed Tests (8xH100-8xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_8 + num_gpus: 8 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - examples/offline_inference/torchrun_dp_example.py + - vllm/config/parallel.py - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/distributed/ - - tests/v1/shutdown - - tests/v1/worker/test_worker_memory_snapshot.py + - vllm/v1/engine/llm_engine.py + - vllm/v1/executor/uniproc_executor.py + - vllm/v1/worker/gpu_worker.py - vllm/platforms/rocm.py commands: - export TORCH_NCCL_BLOCKING_WAIT=1 - - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - - pytest -v -s v1/worker/test_worker_memory_snapshot.py + - 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 --------------------------------------------------------# -- label: Distributed Model Tests (2 GPUs) # TBD +- label: Entrypoints Integration (API Server 2) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/model_loader/sharded_state_loader.py - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/basic_correctness/ - - tests/model_executor/model_loader/test_sharded_state_loader.py - - tests/models/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/entrypoints/rpc + - tests/entrypoints/serve/instrumentator + - tests/tool_use commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/serve/instrumentator + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc + - pytest -v -s tool_use +- label: Entrypoints Integration (API Server openai - Part 1) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py -- label: Plugin Tests (2 GPUs) # TBD +- label: Entrypoints Integration (API Server openai - Part 2) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/plugins/ - - tests/plugins/ - - vllm/platforms/rocm.py - commands: - # BEGIN: platform plugin and general plugin tests, all the code in-between runs on dummy platform - - pip install -e ./plugins/vllm_add_dummy_platform - - pytest -v -s plugins_tests/test_platform_plugins.py - - pip uninstall vllm_add_dummy_platform -y - # END: platform plugin tests - # BEGIN: `io_processor` plugins test, all the code in between uses the `prithvi_io_processor` plugin - - pip install -e ./plugins/prithvi_io_processor_plugin - - pytest -v -s plugins_tests/test_io_processor_plugins.py - - pip uninstall prithvi_io_processor_plugin -y - # END: `io_processor` plugins test - # BEGIN: `bge_m3_sparse io_processor` test - - pip install -e ./plugins/bge_m3_sparse_plugin - - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py - - pip uninstall bge_m3_sparse_plugin -y - # END: `bge_m3_sparse io_processor` test - # BEGIN: `stat_logger` plugins test - - pip install -e ./plugins/vllm_add_dummy_stat_logger - - pytest -v -s plugins_tests/test_stats_logger_plugins.py - - pip uninstall dummy_stat_logger -y - # END: `stat_logger` plugins test - # BEGIN: other tests - - pytest -v -s plugins_tests/test_scheduler_plugins.py - - pip install -e ./plugins/vllm_add_dummy_model - - pytest -v -s distributed/test_distributed_oot.py - - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py - - pytest -v -s models/test_oot_registration.py - - pytest -v -s plugins/lora_resolvers - # END: other tests - - -- label: Pipeline + Context Parallelism (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s distributed/test_pp_cudagraph.py - - pytest -v -s distributed/test_pipeline_parallel.py - - -- label: Ray Dependency Compatibility Check # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/" - source_file_dependencies: - - requirements/ - - setup.py - - vllm/platforms/rocm.py - commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh - - -- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh - - -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - -- label: Distributed Tests (2 GPUs)(H100-MI250) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi325_2 - num_gpus: 2 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/distributed/test_context_parallel.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/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 - - -##################################################################################################################################### -# # -# gfx942 # -# # -##################################################################################################################################### - - -- label: Entrypoints Integration (LLM) # 13.1m - timeout_in_minutes: 22 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/llm - - tests/entrypoints/offline_mode - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py - - pytest -v -s entrypoints/llm/test_generate.py - - pytest -v -s entrypoints/offline_mode - - -- label: Entrypoints Integration (API Server openai - Part 1) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - - -- label: Entrypoints Integration (API Server openai - Part 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - fast_check: true - torch_nightly: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils + - vllm/ + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai/speech_to_text/ - pytest -v -s entrypoints/test_chat_utils.py - - label: Entrypoints Integration (API Server openai - Part 3) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" @@ -1498,33 +1342,30 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - + - 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) #26.9m - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Entrypoints Integration (LLM) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - tests/entrypoints/llm + - tests/entrypoints/offline_mode commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use - + - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py + - pytest -v -s entrypoints/llm/test_generate.py + - pytest -v -s entrypoints/offline_mode -- label: Entrypoints Integration (Pooling) # 22.8m - timeout_in_minutes: 48 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Entrypoints Integration (Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" @@ -1535,446 +1376,260 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling - -- label: Distributed Torchrun + Examples (4 GPUs) # TBD - timeout_in_minutes: 80 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - tests/distributed/test_torchrun_example.py - - tests/distributed/test_torchrun_example_moe.py - - examples/rl/ - - tests/examples/offline_inference/data_parallel.py - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py - - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py - - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py - - 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/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) # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - tests/v1/distributed - - tests/v1/engine/test_engine_core_client.py - - tests/distributed/test_utils - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py - - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py - - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - - pytest -v -s distributed/test_utils.py - - -- label: Distributed Compile + Comm (4 GPUs) # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - tests/distributed/test_pynccl - - tests/distributed/test_events - - tests/compile/fullgraph/test_basic_correctness.py - - tests/distributed/test_symm_mem_allreduce.py - - tests/distributed/test_multiproc_executor.py - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s compile/fullgraph/test_basic_correctness.py - - pytest -v -s distributed/test_pynccl.py - - pytest -v -s distributed/test_events.py - - pytest -v -s distributed/test_symm_mem_allreduce.py - - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node - - -- label: Distributed Tests (8 GPUs)(H100-MI325) # 6.4m - timeout_in_minutes: 10 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_8 - num_gpus: 8 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - 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 - - vllm/platforms/rocm.py - commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep - - -- label: Elastic EP Scaling Test # TBD +- label: Entrypoints Integration (Responses API) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/compilation/ - - tests/distributed/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s distributed/test_elastic_ep.py - - -- label: Engine # 11.3m - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true + torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/engine - - tests/test_sequence - - tests/test_config - - tests/test_logger - - tests/test_vllm_port - commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py - - -- label: Engine (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/engine/ - - tests/v1/engine/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/engine/test_preprocess_error_handling.py - - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py - - -- label: e2e Scheduling (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/general/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general/test_async_scheduling.py - - -- label: e2e Core (1 GPU) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/ - - tests/v1/e2e/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py - - -- label: Spec Decode Eagle # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" - - -- label: Spec Decode Speculators + MTP # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py + - tests/entrypoints/openai/responses commands: - - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" - + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/responses -- label: Spec Decode Ngram + Suffix # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true +- label: Entrypoints Unit Tests # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ + - vllm/entrypoints + - tests/entrypoints/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" - + - pytest -v -s entrypoints/openai/tool_parsers + - 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: Spec Decode Draft Model # TBD +- label: OpenAI API correctness # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ + - csrc/ + - vllm/entrypoints/openai/ + - vllm/model_executor/models/whisper.py - vllm/model_executor/layers/ - - tests/v1/e2e/spec_decode/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ commands: - - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + - bash ../tools/install_torchcodec_rocm.sh || exit 1 + - pytest -s entrypoints/openai/correctness/ +#----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# -- label: V1 e2e (2 GPUs) # 7.1m - timeout_in_minutes: 12 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - optional: true - working_dir: "/vllm-workspace/tests" +- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + num_gpus: 1 + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/ - - tests/v1/e2e + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/backends/mla/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" - + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 -- label: V1 e2e (4 GPUs) # 52.6m - timeout_in_minutes: 106 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true +- label: LM Eval Small Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/e2e + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" - + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt -- label: V1 e2e (4xH100-4xMI325) # TBD +- label: LM Eval Small Models (MI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true - source_file_dependencies: - - vllm/v1/attention/backends/utils.py - - vllm/v1/worker/gpu_model_runner.py - - tests/v1/e2e/test_hybrid_chunked_prefill.py - commands: - - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py - - -- label: V1 Spec Decode # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/ - - tests/v1/spec_decode + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s -m 'not slow_test' v1/spec_decode - + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: GPQA Eval (GPT-OSS) (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/evals/gpt_oss/ commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true +- label: LM Eval Small Models (2xB200-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - # - export HSA_NO_SCRATCH_RECLAIM=1 - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: V1 Speculative Decoding (slow) # TBD +- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/v1/spec_decode/ - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/backends/mla/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: Acceptance Length Test (Large Models) # TBD +- label: LM Eval Large Models (4xA100-4xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/mlp_speculator.py - - tests/v1/spec_decode/test_acceptance_length.py + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test - + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 -- label: V1 attention (H100-MI325) # 14.5m - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" +- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + optional: true + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ - vllm/_aiter_ops.py - - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention - + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Batch Invariance (H100-MI325) # 5.2m - timeout_in_minutes: 12 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/v1/attention - - vllm/model_executor/layers - - tests/v1/determinism/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/spec_decode/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pip install pytest-timeout pytest-forked - - pytest -v -s v1/determinism/test_batch_invariance.py - - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - + - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 -- label: V1 others (CPU) # 10.4m - timeout_in_minutes: 28 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - no_gpu: true +- label: LM Eval Large Models (8xH200-8xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_8 optional: true + num_gpus: 8 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1 + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/layernorm.py + - csrc/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/evals/ commands: - - pytest -v -s -m 'cpu_test' v1/core - - pytest -v -s v1/structured_output - - pytest -v -s v1/test_serial_utils.py - - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'cpu_test' v1/metrics + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt +#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------# -- label: Examples # 24.5m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Examples # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/examples" source_file_dependencies: @@ -2005,54 +1660,12 @@ steps: - 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 ----------------------------------------------------------# -- label: Platform Tests (CUDA) # 5.0m - timeout_in_minutes: 9 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/cuda - commands: - - pytest -v -s cuda/test_cuda_context.py - - pytest -v -s cuda/test_platform_no_cuda_init.py - - -- label: PyTorch Compilation Passes Unit Tests # TBD +- label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/compile/passes - commands: - - pytest -s -v compile/passes --ignore compile/passes/distributed - - -- label: Kernels Core Operation Test # 26.8m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - tests/kernels/core - - tests/kernels/test_top_k_per_row.py - - tests/kernels/test_concat_mla_q.py - - vllm/model_executor/layers/rotary_embedding/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/core kernels/test_top_k_per_row.py - - -- label: Kernels Attention Test %N # 17.7m - timeout_in_minutes: 28 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2066,29 +1679,26 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: Kernels Quantization Test %N # 15.2m - timeout_in_minutes: 24 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - parallelism: 2 +- label: Kernels Core Operation Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/layers/quantization - - tests/kernels/quantization - - tests/kernels/quantization/test_rocm_skinny_gemms.py + - csrc/ + - tests/kernels/core + - tests/kernels/test_top_k_per_row.py + - tests/kernels/test_concat_mla_q.py + - vllm/model_executor/layers/rotary_embedding/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - vllm/model_executor/kernels/ commands: - - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py - label: Kernels MoE Test %N # TBD - timeout_in_minutes: 19 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 parallelism: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2105,11 +1715,27 @@ steps: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +- label: Kernels Quantization Test %N # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + parallelism: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/ + - vllm/model_executor/layers/quantization + - tests/kernels/quantization + - tests/kernels/quantization/test_rocm_skinny_gemms.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/model_executor/kernels/ + commands: + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels FP8 MoE Test # TBD +- label: Kernels FP8 MoE Test (2xH100-2xMI300) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2123,40 +1749,217 @@ steps: commands: - pytest -v -s kernels/moe/test_deepep_moe.py +#----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# -- label: ROCm AITER Ops Test # TBD +- label: LoRA TP (Distributed) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/_aiter_ops.py - - vllm/envs.py + - vllm/lora + - tests/lora - vllm/platforms/rocm.py - - tests/rocm/aiter/ - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py commands: - - pytest -v -s rocm/aiter/ + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + - pytest -v -s -x lora/test_chatglm3_tp.py + - pytest -v -s -x lora/test_llama_tp.py + - pytest -v -s -x lora/test_llm_with_multi_loras.py + - pytest -v -s -x lora/test_olmoe_tp.py + - pytest -v -s -x lora/test_gptoss_tp.py + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py +#---------------------------------------------------------- mi300 · models -----------------------------------------------------------# -- label: Benchmarks # 8.2m - timeout_in_minutes: 20 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Language Models Test (Extended Pooling) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/language/pooling + commands: + - pytest -v -s models/language/pooling -m 'not core_model' + +- label: Language Models Tests (Standard) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true - working_dir: "/vllm-workspace/.buildkite" + torch_nightly: true + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - benchmarks/ + - vllm/ + - tests/models/language + commands: + - pip freeze | grep -E 'torch' + - pytest -v -s models/language -m 'core_model and (not slow_test)' + +- label: Multi-Modal Models (Extended Generation 1) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + - tests/models/multimodal/test_mapping.py + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py + - pytest -v -s models/multimodal/test_mapping.py + +- label: Multi-Modal Models (Extended Generation 2) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' + +- label: Multi-Modal Models (Extended Generation 3) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' + +- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" + - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model + +- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + torch_nightly: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + - tests/models/multimodal/test_mapping.py + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" + - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model + +- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation + - tests/models/multimodal/test_mapping.py + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model + - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model + +- label: Multi-Modal Processor # 1h 42m + timeout_in_minutes: 138 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + - tests/models/registry.py + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/processing/test_tensor_schema.py + +- label: Multi-Modal Processor (CPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + no_gpu: true + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal + - tests/models/registry.py + commands: + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py + +- label: Quantized Models Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/model_executor/layers/quantization + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - tests/models/quantization + - vllm/model_executor/model_loader/ + commands: + - pytest -v -s models/quantization + +- label: Transformers Nightly Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/multimodal/ + - vllm/model_executor/layers/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - tests/models/ + - examples/ commands: - - bash scripts/run-benchmarks.sh + - pip install --upgrade git+https://github.com/huggingface/transformers + - pytest -v -s tests/models/test_initialization.py + - pytest -v -s tests/models/test_transformers.py + - pytest -v -s tests/models/multimodal/processing/ + - pytest -v -s tests/models/multimodal/test_mapping.py + - python3 examples/basic/offline_inference/chat.py + - 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 --------------------------------------------------------# -- label: Quantization # 36.1m - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Quantization # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/ @@ -2177,261 +1980,293 @@ steps: - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py +#----------------------------------------------------------- mi300 · rocm ------------------------------------------------------------# -- label: Language Models Tests (Standard) # 22.8m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - torch_nightly: true +- label: ROCm AITER Ops Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/language + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + - tests/rocm/aiter/ + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py commands: - - pip freeze | grep -E 'torch' - - pytest -v -s models/language -m 'core_model and (not slow_test)' + - pytest -v -s rocm/aiter/ +#--------------------------------------------------------- mi300 · samplers ----------------------------------------------------------# -- label: Language Models Tests (Hybrid) %N # 34.9m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true - parallelism: 2 +- label: Samplers Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/language/generation + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - vllm/v1/sample/ + - vllm/beam_search.py + - tests/samplers + - tests/conftest.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' - - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s samplers +#------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# -- label: Language Models Test (Extended Generation) # 32.2m - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Python-only Installation # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/language/generation + - tests/standalone_tests/python_only_compile.sh + - setup.py + - vllm/platforms/rocm.py commands: - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' - - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - + - bash standalone_tests/python_only_compile.sh -- label: Multi-Modal Processor # 1h 42m - timeout_in_minutes: 138 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true +- label: Regression # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal - - tests/models/registry.py + - tests/test_regression commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/processing/test_tensor_schema.py + - pip install modelscope + - pytest -v -s test_regression.py +#--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD +- label: Ray Dependency Compatibility Check # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/" source_file_dependencies: - - vllm/ - - tests/models/multimodal + - requirements/ + - setup.py + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model + - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh +#------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD +- label: Acceptance Length Test (Large Models) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal + - vllm/v1/spec_decode/ + - vllm/model_executor/models/mlp_speculator.py + - tests/v1/spec_decode/test_acceptance_length.py + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD +- label: e2e Core (1 GPU) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py + - vllm/v1/ + - tests/v1/e2e/ + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model - + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD +- label: e2e Scheduling (1 GPU) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - torch_nightly: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - + - pytest -v -s v1/e2e/general/test_async_scheduling.py -- label: Multi-Modal Models (Extended Generation 1) # 1h 2m - timeout_in_minutes: 106 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: Engine (1 GPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py + - vllm/v1/engine/ + - tests/v1/engine/ + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py - + - pytest -v -s v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py -- label: Multi-Modal Models (Extended Generation 2) # TBD +- label: Spec Decode Draft Model # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' + - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" +- label: Spec Decode Eagle # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" -- label: Multi-Modal Models (Extended Generation 3) # TBD +- label: Spec Decode Ngram + Suffix # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - + - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" -- label: Multi-Modal Models (Extended Pooling) # TBD +- label: Spec Decode Speculators + MTP # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - + - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" -- label: Quantized Models Test # 21.4m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: V1 attention (H100-MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/layers/quantization + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention - vllm/_aiter_ops.py + - vllm/envs.py - vllm/platforms/rocm.py - - tests/models/quantization - - vllm/model_executor/model_loader/ commands: - - pytest -v -s models/quantization - + - pytest -v -s v1/attention + +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + # - export HSA_NO_SCRATCH_RECLAIM=1 + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: Transformers Nightly Models # 50.9m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 +- label: V1 others (CPU) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + no_gpu: true optional: true - working_dir: "/vllm-workspace/" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/multimodal/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/models/ - - examples/ + - vllm/ + - tests/v1 commands: - - pip install --upgrade git+https://github.com/huggingface/transformers - - pytest -v -s tests/models/test_initialization.py - - pytest -v -s tests/models/test_transformers.py - - pytest -v -s tests/models/multimodal/processing/ - - pytest -v -s tests/models/multimodal/test_mapping.py - - python3 examples/basic/offline_inference/chat.py - - 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 - + - pytest -v -s -m 'cpu_test' v1/core + - pytest -v -s v1/structured_output + - pytest -v -s v1/test_serial_utils.py + - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/metrics -- label: Quantized MoE Test (B200-MI325) # TBD +- label: V1 Sample + Logits # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - tests/quantization/test_gfx3xx_moe.py - - vllm/model_executor/models/deepseek_v2.py - - vllm/model_executor/models/gpt_oss.py - - vllm/model_executor/models/llama4.py - - vllm/model_executor/layers/fused_moe - - vllm/model_executor/layers/quantization/compressed_tensors - - vllm/model_executor/layers/quantization/modelopt.py - - vllm/model_executor/layers/quantization/mxfp4.py - - vllm/v1/attention/backends/triton_attn.py - - vllm/v1/attention/backends/rocm_attn.py - - vllm/v1/attention/backends/rocm_aiter_fa.py - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py commands: - - pytest -s -v tests/quantization/test_gfx3xx_moe.py - + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py -- label: Distributed DP Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 +- label: Distributed DP Tests (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2451,152 +2286,126 @@ steps: - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py - -- label: Distributed Compile + RPC Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 +- label: Distributed Tests (2xH100-2xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 num_gpus: 2 - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/" source_file_dependencies: - - vllm/compilation/ - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/compile/fullgraph/test_basic_correctness.py - - tests/compile/test_wrapper.py - - tests/entrypoints/llm/test_collective_rpc.py + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - tests/v1/distributed/test_dbo.py + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s entrypoints/llm/test_collective_rpc.py - - pytest -v -s ./compile/fullgraph/test_basic_correctness.py - - pytest -v -s ./compile/test_wrapper.py - + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py + - 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 Torchrun + Shutdown Tests (2 GPUs) # 56.1m - timeout_in_minutes: 102 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 +- label: Metrics, Tracing (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/distributed/ - - tests/v1/shutdown - - tests/v1/worker/test_worker_memory_snapshot.py - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/tracing commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - - pytest -v -s v1/worker/test_worker_memory_snapshot.py - + - "pip install \ + 'opentelemetry-sdk>=1.26.0' \ + 'opentelemetry-api>=1.26.0' \ + 'opentelemetry-exporter-otlp>=1.26.0' \ + 'opentelemetry-semantic-conventions-ai>=0.4.1'" + - pytest -v -s v1/tracing -- label: Distributed Model Tests (2 GPUs) # 19.3m - timeout_in_minutes: 38 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 +- label: V1 e2e (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/model_loader/sharded_state_loader.py - - vllm/model_executor/models/ - - vllm/model_executor/layers/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/basic_correctness/ - - tests/model_executor/model_loader/test_sharded_state_loader.py - - tests/models/ + - vllm/ + - tests/v1/e2e commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" -- label: LoRA TP (Distributed) # 9.8m - timeout_in_minutes: 18 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 +- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/lora - - tests/lora + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True - - pytest -v -s -x lora/test_chatglm3_tp.py - - pytest -v -s -x lora/test_llama_tp.py - - pytest -v -s -x lora/test_llm_with_multi_loras.py - - pytest -v -s -x lora/test_olmoe_tp.py - - pytest -v -s -x lora/test_gptoss_tp.py - - pytest -v -s -x lora/test_qwen35_densemodel_lora.py - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Weight Loading Multiple GPU # 7.5m - timeout_in_minutes: 14 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 +- label: Distributed DP Tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/weight_loading + - vllm/distributed/ + - tests/v1/distributed + - tests/v1/engine/test_engine_core_client.py + - tests/distributed/test_utils + - vllm/platforms/rocm.py commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py + - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp + - pytest -v -s distributed/test_utils.py -- label: Weight Loading Multiple GPU - Large Models # 12.6m - timeout_in_minutes: 26 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/weight_loading + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Ray Dependency Compatibility Check # TBD +- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/" + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/ - - setup.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: Distributed NixlConnector PD accuracy (4 GPUs) # 27.4m - timeout_in_minutes: 44 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 +- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 num_gpus: 4 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -2604,161 +2413,160 @@ steps: - vllm/platforms/rocm.py commands: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: V1 e2e (4 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/e2e + commands: + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy" + +- label: V1 e2e (4xH100-4xMI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + optional: true + source_file_dependencies: + - vllm/v1/attention/backends/utils.py + - vllm/v1/worker/gpu_model_runner.py + - tests/v1/e2e/test_hybrid_chunked_prefill.py + commands: + - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py +#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------# -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: Weight Loading Multiple GPU # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py + - vllm/ + - tests/weight_loading commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt -- label: Distributed Tests (4 GPUs)(A100-MI325) # 20.9m - timeout_in_minutes: 37 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 +- label: Weight Loading Multiple GPU - Large Models # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ + - tests/weight_loading commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s distributed/test_custom_all_reduce.py - - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - pytest -v -s -x lora/test_mixtral.py + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt + +######################################################################################################################################### +# # +# MI325 (gfx942) tests # +# # +######################################################################################################################################### +#---------------------------------------------------------- mi325 · compile ----------------------------------------------------------# -- label: Distributed Tests (2 GPUs)(H100-MI325) # TBD +- label: Distributed Compile + RPC Tests (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_2 num_gpus: 2 - working_dir: "/vllm-workspace/" + working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/compilation/ - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - tests/v1/distributed/test_dbo.py - - vllm/_aiter_ops.py + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/compile/fullgraph/test_basic_correctness.py + - tests/compile/test_wrapper.py + - tests/entrypoints/llm/test_collective_rpc.py - vllm/platforms/rocm.py commands: - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s entrypoints/llm/test_collective_rpc.py + - pytest -v -s ./compile/fullgraph/test_basic_correctness.py + - pytest -v -s ./compile/test_wrapper.py +#-------------------------------------------------------- mi325 · distributed --------------------------------------------------------# -- label: Distributed Compile Unit Tests (2xH100-2xMI325) # 14.3m - timeout_in_minutes: 32 +- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_2 num_gpus: 2 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers - - tests/compile/passes/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py - # TODO: this test is not supported on ROCm, there are aiter kernels for this. - # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py - # - pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm - # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" - - -- label: LM Eval Small Models # 13.3m - timeout_in_minutes: 23 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt - - -- label: LM Eval Small Models (MI325) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py - vllm/platforms/rocm.py commands: - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown + - pytest -v -s v1/worker/test_worker_memory_snapshot.py -- label: LM Eval Small Models (B200-MI325) # TBD +- label: Distributed Compile + Comm (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 + agent_pool: mi325_4 + num_gpus: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/ + - tests/distributed/test_pynccl + - tests/distributed/test_events + - tests/compile/fullgraph/test_basic_correctness.py + - tests/distributed/test_symm_mem_allreduce.py + - tests/distributed/test_multiproc_executor.py - vllm/platforms/rocm.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s compile/fullgraph/test_basic_correctness.py + - pytest -v -s distributed/test_pynccl.py + - pytest -v -s distributed/test_events.py + - pytest -v -s distributed/test_symm_mem_allreduce.py + - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node +#---------------------------------------------------------- mi325 · engine -----------------------------------------------------------# -- label: LM Eval Large Models (H200-MI325) # TBD +- label: Engine # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_8 - optional: true - num_gpus: 8 + agent_pool: mi325_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py - - csrc/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/ + - vllm/ + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py +#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------# -- label: LM Eval Large Models (4 GPUs)(FP8) # 24.8m - timeout_in_minutes: 42 +- label: LM Eval Large Models (4xH100-4xMI325) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_4 num_gpus: 4 @@ -2777,29 +2585,7 @@ steps: - export VLLM_USE_DEEP_GEMM=0 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 - -- label: LM Eval Large Models (4 GPUs)(A100-MI325) # 17.3m - timeout_in_minutes: 27 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 - - -- label: ROCm LM Eval Large Models (8 Card) # TBD +- label: ROCm LM Eval Large Models (8 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_8 @@ -2820,232 +2606,143 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 +#---------------------------------------------------------- mi325 · models -----------------------------------------------------------# -- label: GPQA Eval (GPT-OSS) (H100-MI325) # TBD +- label: Language Models Test (Extended Generation) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_2 - num_gpus: 2 - optional: true + agent_pool: mi325_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/fused_moe/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/gpt_oss/ + - vllm/ + - tests/models/language/generation commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt - + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' -- label: DeepSeek V2-Lite Accuracy # 6.7m - timeout_in_minutes: 12 +- label: Language Models Tests (Hybrid) %N # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" + agent_pool: mi325_1 + torch_nightly: true + parallelism: 2 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/models/language/generation commands: - - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 - + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' + - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100-MI325) # TBD +- label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_1 - num_gpus: 1 - working_dir: "/vllm-workspace" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/models/multimodal/pooling commands: - - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 - - -- label: Qwen3-30B-A3B-FP8-block Accuracy # 6.4m - timeout_in_minutes: 11 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - pytest -v -s models/multimodal/pooling -m 'not core_model' + +- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + torch_nightly: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" + - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model +#------------------------------------------------------------ mi325 · v1 -------------------------------------------------------------# -- label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy # 10.9m - timeout_in_minutes: 22 +- label: V1 Spec Decode # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" + agent_pool: mi325_1 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/spec_decode/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/spec_decode commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 - -##### .buildkite/test_areas/compile.yaml ##### -# Slowly setting up the tests so that it is also easier for the -# CI team to review and upstream to the pipelinev2. -# The following tests are important for vLLM IR Ops refactoring, -# which affects fusion passes on ROCm. So we have to -# enable them as as soon as possible. + - pytest -v -s -m 'not slow_test' v1/spec_decode -## TODO: Enable the test in this group -# # corresponds to .buildkite/test_areas/compile.yaml -# - label: Fusion and Compile Unit Tests (2xB200-2xMI325) # TBD -# timeout_in_minutes: 180 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325, tj] -# agent_pool: mi325_1 # changed to 1 GPU until the fusion all reduce is enabled then only revert back to 2 GPUs -# num_gpus: 1 -# working_dir: "/vllm-workspace/" -# source_file_dependencies: -# - csrc/quantization/fp4/ -# - vllm/model_executor/layers/quantization/ -# - vllm/model_executor/layers/layernorm.py -# - vllm/model_executor/layers/activation.py -# - vllm/model_executor/layers/attention/attention.py -# - vllm/v1/attention/backends/flashinfer.py -# - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes -# - tests/compile/test_fusion_attn.py -# - tests/compile/test_silu_mul_quant_fusion.py -# - tests/compile/distributed/test_fusion_all_reduce.py -# - tests/compile/fullgraph/test_full_graph.py -# commands: -# - rocm-smi -# # we run all backend tests on ROCm -# # These two tests are covered in "PyTorch Compilation Passes Unit Tests" -# # - "pytest -v -s tests/compile/passes/test_fusion_attn.py" -# # - "pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py" -# # TODO: this test is not supported on ROCm, there are aiter kernels for this. -# # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py -# # TODO: find out more details -# # - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile +######################################################################################################################################### +# # +# MI355 (gfx950) tests # +# # +######################################################################################################################################### +#-------------------------------------------------------- mi355 · benchmarks ---------------------------------------------------------# -- label: Fusion E2E Quick (H100-MI325) # TBD +- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - num_gpus: 1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + num_gpus: 2 working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/ + - benchmarks/attention_benchmarks/ - vllm/v1/attention/ - - vllm/compilation/ - - tests/compile/fusions_e2e/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - rocm-smi - # Run all models and attn backends but only Inductor partition and native custom ops - - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'" - # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER - - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'" + - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 +#-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# -- label: Fusion E2E Config Sweep (H100-MI325) # TBD +- label: Distributed Tests (2xH100-2xMI355) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - num_gpus: 1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/quantization/ - - vllm/compilation/ - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/attention/attention.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/fusions_e2e/ + - vllm/distributed/ + - vllm/v1/distributed/ + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - tests/distributed/test_context_parallel.py + - tests/v1/distributed/test_dbo.py + - examples/offline_inference/data_parallel.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - rocm-smi - - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" - -## There are no ops on ROCm for these tests. -## The test still passes but the logs are not useful. -## fused ops just call torch.ops.symm_mem which -## exists in ROCm even though they don't work -# - label: AsyncTP Correctness Tests (2xH100-2xMI325) -# - label: Fusion E2E TP2 Quick (H100-MI325) -# - label: Fusion E2E TP2 AsyncTP Config Sweep (H100-MI325) -# - label: Fusion E2E TP2 (B200-MI325) -# - label: Sequence Parallel Correctness Tests (2xH100-2xMI325) - + - export TORCH_NCCL_BLOCKING_WAIT=1 + - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/distributed/test_dbo.py -##################################################################################################################################### -# # -# gfx950 # -# # -##################################################################################################################################### +#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# -- label: Entrypoints Integration (API Server openai - Part 1) # TBD +- label: Entrypoints Integration (API Server 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 + optional: true fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/openai - - tests/entrypoints/test_chat_utils + - tests/entrypoints/rpc + - tests/entrypoints/serve/instrumentator + - tests/tool_use commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - + - pytest -v -s entrypoints/serve/instrumentator + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc + - pytest -v -s tool_use -- label: Entrypoints Integration (API Server openai - Part 2) # TBD +- label: Entrypoints Integration (API Server openai - Part 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 @@ -3058,12 +2755,9 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai/speech_to_text/ - - pytest -v -s entrypoints/test_chat_utils.py - + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py -- label: Entrypoints Integration (API Server openai - Part 3) # TBD +- label: Entrypoints Integration (API Server openai - Part 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 @@ -3076,28 +2770,24 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py + - pytest -v -s entrypoints/openai/speech_to_text/ + - pytest -v -s entrypoints/test_chat_utils.py -- label: Entrypoints Integration (API Server 2) # TBD +- label: Entrypoints Integration (API Server openai - Part 3) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - optional: true fast_check: true torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - tests/entrypoints/openai + - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use - + - 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 (Pooling) # TBD timeout_in_minutes: 180 @@ -3113,117 +2803,112 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/pooling +#----------------------------------------------------------- mi355 · evals -----------------------------------------------------------# -- label: Regression # TBD +- label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 + agent_pool: mi355_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_regression - commands: - - pip install modelscope - - pytest -v -s test_regression.py - - -- label: V1 Spec Decode # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/spec_decode + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - tests/evals/gpt_oss/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s -m 'not slow_test' v1/spec_decode - + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt -- label: V1 Sample + Logits # TBD - timeout_in_minutes: 60 +- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD + timeout_in_minutes: 120 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 + agent_pool: mi355_2 + num_gpus: 2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/sample - - tests/v1/logits_processors - - tests/v1/test_oracle.py - - tests/v1/test_request.py - - tests/v1/test_outputs.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/transformers_utils/configs/qwen3_5_moe.py + - vllm/model_executor/models/qwen.py + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/model_executor/layers/fla/ops/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - pytest -v -s v1/sample - - pytest -v -s v1/logits_processors - - pytest -v -s v1/test_oracle.py - - pytest -v -s v1/test_request.py - - pytest -v -s v1/test_outputs.py - + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt -- label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 60 +- label: LM Eval Small Models (2xB200-2xMI355) # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 + agent_pool: mi355_2 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s -m 'not cpu_test' v1/core - - pytest -v -s v1/executor - - pytest -v -s v1/kv_offload - - pytest -v -s v1/worker - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api - - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: V1 Speculative Decoding (slow) # TBD +- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" + agent_pool: mi355_2 + num_gpus: 2 + working_dir: "/vllm-workspace" source_file_dependencies: - - vllm/v1/spec_decode/ - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/eplb + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 -- label: V1 attention (B200-MI355) # TBD +- label: LM Eval Large Models (4xH100-4xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py - vllm/_aiter_ops.py - - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - export VLLM_USE_DEEP_GEMM=0 + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 +#--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# - label: Examples # TBD timeout_in_minutes: 180 @@ -3258,6 +2943,31 @@ steps: - 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 ----------------------------------------------------------# + +- label: Kernels (B200-MI355) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_1 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - csrc/quantization/fp4/ + - csrc/attention/mla/ + - csrc/quantization/cutlass_w8a8/moe/ + - vllm/model_executor/layers/fused_moe/cutlass_moe.py + - vllm/v1/attention/backends/triton_attn.py + - vllm/v1/attention/backends/rocm_attn.py + - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/v1/attention/backends/mla/aiter_triton_mla.py + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + commands: + - rocm-smi + - python3 examples/basic/offline_inference/chat.py + - pytest -v -s tests/kernels/attention/test_attention_selector.py - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 @@ -3277,25 +2987,6 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: Kernels Quantization Test %N # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - parallelism: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/quantization/ - - vllm/model_executor/layers/quantization - - tests/kernels/quantization - - tests/kernels/quantization/test_rocm_skinny_gemms.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/kernels/ - commands: - - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - - label: Kernels MoE Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3316,54 +3007,40 @@ steps: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: Kernels FP8 MoE Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/moe/ - - csrc/quantization/w8a8/cutlass/moe/ - - vllm/model_executor/layers/fused_moe/ - - tests/kernels/moe/test_deepep_moe.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/envs.py - commands: - - pytest -v -s kernels/moe/test_deepep_moe.py - - -- label: Quantization # TBD +- label: Kernels Quantization Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 + parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ + - csrc/quantization/ - vllm/model_executor/layers/quantization - - tests/quantization + - tests/kernels/quantization + - tests/kernels/quantization/test_rocm_skinny_gemms.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/model_executor/kernels/ commands: - - uv pip install --system torchao==0.17.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py - + - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Language Models Tests (Standard) # TBD +- label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true + agent_pool: mi355_2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/language + - csrc/moe/ + - csrc/quantization/w8a8/cutlass/moe/ + - vllm/model_executor/layers/fused_moe/ + - tests/kernels/moe/test_deepep_moe.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - vllm/envs.py commands: - - pip freeze | grep -E 'torch' - - pytest -v -s models/language -m 'core_model and (not slow_test)' + - pytest -v -s kernels/moe/test_deepep_moe.py +#---------------------------------------------------------- mi355 · models -----------------------------------------------------------# - label: Language Models Test (Extended Generation) # TBD timeout_in_minutes: 180 @@ -3375,10 +3052,9 @@ steps: - tests/models/language/generation commands: - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - - label: Language Models Test (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3390,71 +3066,46 @@ steps: commands: - pytest -v -s models/language/pooling -m 'not core_model' - -- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" - - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model - - -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - torch_nightly: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model - - -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD +- label: Language Models Test (PPL) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - torch_nightly: true - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/transformers_utils/configs/qwen3_5.py + - vllm/transformers_utils/configs/qwen3_5_moe.py + - vllm/model_executor/models/qwen.py + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_next.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/model_executor/layers/fla/ops/ + - vllm/_aiter_ops.py + - vllm/v1/attention/backends/triton_attn.py + - vllm/v1/attention/backends/rocm_attn.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/v1/attention/backends/rocm_aiter_fa.py + - vllm/v1/attention/backends/flex_attention.py + - vllm/v1/attention/ops/ + - vllm/platforms/rocm.py + - tests/models/language/generation_ppl_test commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" - - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model - + - pytest -v -s models/language/generation_ppl_test -- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD +- label: Language Models Tests (Standard) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 torch_nightly: true - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/models/multimodal/generation + - tests/models/language commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing - - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model - + - pip freeze | grep -E 'torch' + - pytest -v -s models/language -m 'core_model and (not slow_test)' - label: Multi-Modal Models (Extended Generation 1) # TBD timeout_in_minutes: 180 @@ -3471,21 +3122,6 @@ steps: - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py - -- label: Multi-Modal Models (Extended Generation 2) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' - - - label: Multi-Modal Models (Extended Generation 3) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3499,7 +3135,6 @@ steps: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' - - label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3512,303 +3147,252 @@ steps: commands: - pytest -v -s models/multimodal/pooling -m 'not core_model' - -- label: Quantized Models Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/models/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/model_loader/ - commands: - - pytest -v -s models/quantization - - -- label: Kernels (B200-MI355) # TBD +- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - working_dir: "/vllm-workspace/" - source_file_dependencies: - - csrc/quantization/fp4/ - - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ - - vllm/model_executor/layers/fused_moe/cutlass_moe.py - - vllm/v1/attention/backends/triton_attn.py - - vllm/v1/attention/backends/rocm_attn.py - - vllm/v1/attention/backends/rocm_aiter_fa.py - - vllm/v1/attention/backends/rocm_aiter_unified_attn.py - - vllm/v1/attention/backends/mla/aiter_triton_mla.py - - vllm/v1/attention/backends/mla/rocm_aiter_mla.py - - vllm/v1/attention/selector.py - - vllm/platforms/rocm.py - - vllm/_aiter_ops.py - commands: - - rocm-smi - - python3 examples/basic/offline_inference/chat.py - - pytest -v -s tests/kernels/attention/test_attention_selector.py - - -- label: Weight Loading Multiple GPU # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 + torch_nightly: true + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/weight_loading + - tests/models/multimodal commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" + - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model -- label: Weight Loading Multiple GPU - Large Models # TBD +- label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - working_dir: "/vllm-workspace/tests" - num_gpus: 2 + agent_pool: mi355_1 + torch_nightly: true optional: true + working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/weight_loading + - tests/models/multimodal/generation commands: - - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt - + - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model + - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model -- label: Ray Dependency Compatibility Check # TBD +- label: Quantized Models Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - optional: true - working_dir: "/" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/ - - setup.py + - vllm/model_executor/layers/quantization + - tests/models/quantization + - vllm/_aiter_ops.py - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ commands: - - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh + - pytest -v -s models/quantization +#------------------------------------------------------- mi355 · quantization --------------------------------------------------------# -- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD +- label: Quantization # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_4 - num_gpus: 4 - optional: true + agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ + - csrc/ + - vllm/model_executor/layers/quantization + - tests/quantization + - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - uv pip install --system torchao==0.17.0 + - uv pip install --system conch-triton-kernels + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py + +# - label: Quantized MoE Test (B200-MI355) # TBD +# timeout_in_minutes: 180 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# agent_pool: mi355_1 +# working_dir: "/vllm-workspace/" +# source_file_dependencies: +# - tests/quantization/test_gfx950_moe.py +# - vllm/model_executor/models/deepseek_v2.py +# - vllm/model_executor/models/gpt_oss.py +# - vllm/model_executor/models/llama4.py +# - vllm/model_executor/layers/fused_moe +# - vllm/model_executor/layers/quantization/compressed_tensors +# - vllm/model_executor/layers/quantization/modelopt.py +# - vllm/model_executor/layers/quantization/mxfp4.py +# - vllm/v1/attention/backends/triton_attn.py +# - vllm/v1/attention/backends/rocm_attn.py +# - vllm/v1/attention/backends/rocm_aiter_fa.py +# - vllm/v1/attention/backends/mla/ +# - vllm/v1/attention/selector.py +# - vllm/model_executor/layers/layernorm.py +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# - vllm/model_executor/model_loader/ +# commands: +# - pytest -s -v tests/quantization/test_gfx950_moe.py +#------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD +- label: V1 attention (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_4 - num_gpus: 4 - optional: true + agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ + - vllm/config/attention.py + - vllm/model_executor/layers/attention + - vllm/v1/attention + - tests/v1/attention + - vllm/_aiter_ops.py + - vllm/envs.py - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - + - pytest -v -s v1/attention -- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD - timeout_in_minutes: 180 +- label: V1 Core + KV + Metrics # TBD + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - optional: true + agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py commands: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh - + - pytest -v -s -m 'not cpu_test' v1/core + - pytest -v -s v1/executor + - pytest -v -s v1/kv_offload + - pytest -v -s v1/worker + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'not cpu_test' v1/metrics + - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: Distributed Tests (2 GPUs)(H100-MI355) # TBD - timeout_in_minutes: 180 +- label: V1 Sample + Logits # TBD + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - tests/distributed/test_context_parallel.py - - tests/v1/distributed/test_dbo.py - - examples/offline_inference/data_parallel.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/sample + - tests/v1/logits_processors + - tests/v1/test_oracle.py + - tests/v1/test_request.py + - tests/v1/test_outputs.py commands: - - export TORCH_NCCL_BLOCKING_WAIT=1 - - pytest -v -s tests/distributed/test_context_parallel.py - - 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 - + - pytest -v -s v1/sample + - pytest -v -s v1/logits_processors + - pytest -v -s v1/test_oracle.py + - pytest -v -s v1/test_request.py + - pytest -v -s v1/test_outputs.py -- label: Distributed Compile Unit Tests (2xH100-2xMI355) # TBD - timeout_in_minutes: 180 +- label: V1 Spec Decode # TBD + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" + agent_pool: mi355_1 + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers - - tests/compile/passes/distributed/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/v1/spec_decode commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py - # TODO: this test is not supported on ROCm, there are aiter kernels for this. - # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py - # - pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm - # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'" - + - pytest -v -s -m 'not slow_test' v1/spec_decode -- label: LM Eval Small Models (B200-MI355) # TBD +- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 + num_gpus: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh -- label: LM Eval Large Models (4 GPUs)(FP8) # TBD +- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_4 num_gpus: 4 optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - export VLLM_USE_DEEP_GEMM=0 - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-fp8.txt --tp-size=4 - + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: GPQA Eval (GPT-OSS) (B200-MI355) # TBD +- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx955nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + agent_pool: mi355_4 + num_gpus: 4 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/fused_moe/ - - tests/evals/gpt_oss/ - - vllm/_aiter_ops.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# -- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD +- label: Weight Loading Multiple GPU # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 num_gpus: 2 - working_dir: "/vllm-workspace" + working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/model_executor/layers/fused_moe/ - - vllm/distributed/eplb - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/weight_loading commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 - + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt -- label: Attention Benchmarks Smoke Test (B200-MI355) # TBD +- label: Weight Loading Multiple GPU - Large Models # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 + working_dir: "/vllm-workspace/tests" num_gpus: 2 - working_dir: "/vllm-workspace/" + optional: true source_file_dependencies: - - benchmarks/attention_benchmarks/ - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/weight_loading commands: - - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt +#----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# -- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD - timeout_in_minutes: 120 +- label: Regression # TBD + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_2 - num_gpus: 2 + agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/models/qwen3_5.py - - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - - vllm/model_executor/models/qwen.py - - vllm/model_executor/models/qwen2.py - - vllm/model_executor/models/qwen3.py - - vllm/model_executor/models/qwen3_next.py - - vllm/model_executor/models/qwen3_next_mtp.py - - vllm/model_executor/layers/fla/ops/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py + - vllm/ + - tests/test_regression commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt + - pip install modelscope + - pytest -v -s test_regression.py diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml new file mode 100644 index 000000000000..a10fda41ef0d --- /dev/null +++ b/.buildkite/test_areas/disaggregated.yaml @@ -0,0 +1,98 @@ +group: Disaggregated +depends_on: + - image-build +steps: +- label: Distributed NixlConnector PD accuracy (4 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) + timeout_in_minutes: 20 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh + +- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) + timeout_in_minutes: 30 + device: a100 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh + +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh \ No newline at end of file diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 19815e14b4e0..093f3ab4fe1f 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -196,6 +196,8 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - pytest -v -s tests/distributed/test_packed_tensor.py - label: Distributed Tests (2 GPUs)(B200) device: b200 @@ -224,63 +226,6 @@ steps: commands: - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" -- label: Distributed NixlConnector PD accuracy (4 GPUs) - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs) - timeout_in_minutes: 20 - working_dir: "/vllm-workspace/tests" - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - -- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) - timeout_in_minutes: 30 - device: a100 - working_dir: "/vllm-workspace/tests" - num_devices: 2 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh - - label: Pipeline + Context Parallelism (4 GPUs) timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 5b7f96bc7a26..857fefd268a4 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -29,6 +29,15 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 +- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy + timeout_in_minutes: 60 + device: h100 + optional: true + num_devices: 4 + working_dir: "/vllm-workspace" + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 + - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) timeout_in_minutes: 60 device: h100 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index a05ee886f5b9..2df0a5e253fb 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -18,10 +18,22 @@ steps: source_file_dependencies: - csrc/ - tests/kernels/core - - tests/kernels/test_top_k_per_row.py - tests/kernels/test_concat_mla_q.py commands: - - pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py + +- label: Kernels MiniMax Reduce RMS Test (2 GPUs) + timeout_in_minutes: 15 + num_devices: 2 + device: h100 + source_file_dependencies: + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - vllm/model_executor/layers/mamba/linear_attn.py + - vllm/model_executor/layers/mamba/lamport_workspace.py + - tests/kernels/core/test_minimax_reduce_rms.py + commands: + - pytest -v -s kernels/core/test_minimax_reduce_rms.py - label: Kernels Attention Test %N timeout_in_minutes: 35 @@ -107,6 +119,7 @@ steps: - vllm/v1/attention/backends/mla/flashinfer_mla.py - vllm/v1/attention/selector.py - vllm/platforms/cuda.py + - tests/kernels/test_top_k_per_row.py commands: - nvidia-smi - python3 examples/basic/offline_inference/chat.py @@ -117,6 +130,7 @@ steps: - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py + - pytest -v -s tests/kernels/test_top_k_per_row.py # Quantization - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8' - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py @@ -127,6 +141,7 @@ steps: - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py + - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py - pytest -v -s tests/kernels/moe/test_flashinfer.py - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py @@ -141,7 +156,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==0.3.3 + - pip install helion==1.0.0 - pytest -v -s kernels/helion/ @@ -186,7 +201,14 @@ steps: timeout_in_minutes: 90 device: h100 num_devices: 2 - optional: true + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/device_communicators/ + - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py @@ -195,6 +217,13 @@ steps: timeout_in_minutes: 90 device: b200 num_devices: 2 - optional: true + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/device_communicators/ + - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 39029efe9cd9..a07d702cf3ce 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -91,6 +91,16 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt +- label: LM Eval TurboQuant KV Cache + timeout_in_minutes: 75 + source_file_dependencies: + - vllm/model_executor/layers/quantization/turboquant/ + - vllm/v1/attention/backends/turboquant_attn.py + - vllm/v1/attention/ops/triton_turboquant_decode.py + - vllm/v1/attention/ops/triton_turboquant_store.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt + - label: GPQA Eval (GPT-OSS) (H100) timeout_in_minutes: 120 device: h100 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 82fe6670a990..0cf9ec433925 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -224,6 +224,7 @@ steps: - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - label: Acceptance Length Test (Large Models) # optional timeout_in_minutes: 25 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 8ba4484fb02e..ed782c061fa3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -1,10 +1,9 @@ group: Models - Basic -depends_on: +depends_on: - image-build steps: - label: Basic Models Tests (Initialization) timeout_in_minutes: 45 - device: h200_18gb torch_nightly: true source_file_dependencies: - vllm/ @@ -13,10 +12,11 @@ steps: commands: # Run a subset of model initialization tests - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset + mirror: + torch_nightly: {} - label: Basic Models Tests (Extra Initialization) %N timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/test_initialization.py @@ -27,6 +27,8 @@ steps: # test.) Also run if model initialization test file is modified - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Basic Models Tests (Other) timeout_in_minutes: 45 @@ -42,10 +44,10 @@ steps: device: mi325_1 depends_on: - image-build-amd - + - label: Basic Models Test (Other CPU) # 5min - depends_on: + depends_on: - image-build-cpu timeout_in_minutes: 10 source_file_dependencies: @@ -70,3 +72,18 @@ steps: - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl # Whisper needs spawn method to avoid deadlock - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper + +- label: Transformers Backward Compatibility Models Test + working_dir: "/vllm-workspace/" + optional: true + soft_fail: true + commands: + - pip install transformers==4.57.5 + - pytest -v -s tests/models/test_initialization.py + - pytest -v -s tests/models/test_transformers.py + - pytest -v -s tests/models/multimodal/processing/ + - pytest -v -s tests/models/multimodal/test_mapping.py + - python3 examples/offline_inference/basic/chat.py + - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl + # Whisper needs spawn method to avoid deadlock + - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 7eac9e30193a..c13371e25f1d 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -1,10 +1,9 @@ group: Models - Language -depends_on: +depends_on: - image-build steps: - label: Language Models Tests (Standard) timeout_in_minutes: 25 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language @@ -12,10 +11,11 @@ steps: # Test standard language models, excluding a subset of slow tests - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' + mirror: + torch_nightly: {} - label: Language Models Tests (Extra Standard) %N timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/language/pooling/test_embedding.py @@ -27,10 +27,11 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Tests (Hybrid) %N timeout_in_minutes: 75 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language/generation @@ -42,6 +43,8 @@ steps: # Shard hybrid language model tests - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Test (Extended Generation) # 80min timeout_in_minutes: 110 @@ -62,7 +65,7 @@ steps: - image-build-amd commands: - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: Language Models Test (PPL) diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 3bf907bb6c5c..245ef24026d2 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -28,6 +28,7 @@ steps: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model + - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model mirror: amd: device: mi325_1 @@ -56,7 +57,8 @@ steps: - tests/models/multimodal commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index ad538e919185..a3648219d89d 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -64,6 +64,6 @@ steps: device: h200_18gb soft_fail: true source_file_dependencies: - - requirements/nightly_torch_test.txt + - requirements/test/nightly-torch.txt commands: - bash standalone_tests/pytorch_nightly_dependency.sh diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 0a395ea55883..a42d59b021c6 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -9,7 +9,7 @@ steps: - vllm/model_executor/layers/quantization - tests/quantization commands: - # temporary install here since we need nightly, will move to requirements/test.in + # temporary install here since we need nightly, will move to requirements/test/cuda.in # after torchao 0.12 release, and pin a working version of torchao nightly here # since torchao nightly is only compatible with torch nightly currently diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index a0b730968675..05925da0da01 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -12,6 +12,17 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" +- label: Spec Decode Eagle Nightly B200 + timeout_in_minutes: 30 + device: b200 + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + - label: Spec Decode Speculators + MTP timeout_in_minutes: 30 device: h200_18gb @@ -23,6 +34,18 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" +- label: Spec Decode Speculators + MTP Nightly B200 + timeout_in_minutes: 30 + device: b200 + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" + - label: Spec Decode Ngram + Suffix timeout_in_minutes: 30 device: h200_18gb @@ -42,3 +65,27 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: Spec Decode Draft Model Nightly B200 + timeout_in_minutes: 30 + device: b200 + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: DFlash Speculators Correctness + timeout_in_minutes: 30 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - tests/v1/spec_decode/test_speculators_dflash.py + commands: + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a9c1c17128b8..a20c5e7e9dce 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,7 +3,7 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy -/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery +/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /vllm/pooling_params.py @noooop @DarkLight1337 /vllm/tokenizers @DarkLight1337 @njhill /vllm/renderers @DarkLight1337 @njhill -/vllm/reasoning @aarnphm @chaunceyjiang -/vllm/tool_parsers @aarnphm @chaunceyjiang +/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning # vLLM V1 /vllm/v1/attention @LucasWilkinson @MatthewBonanni @@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson /tests/v1/kv_connector/nixl_integration @NickLucche /tests/v1/kv_connector @ApostaC @orozery /tests/v1/kv_offload @ApostaC @orozery -/tests/v1/determinism @yewentao256 +/tests/v1/determinism @yewentao256 +/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor @@ -120,16 +124,16 @@ mkdocs.yaml @hmellor /tools/pre_commit @hmellor # CPU -/vllm/v1/worker/cpu* @bigPYJ1151 +/vllm/v1/worker/cpu* @bigPYJ1151 @xuechendi /csrc/cpu @bigPYJ1151 -/vllm/platforms/cpu.py @bigPYJ1151 +/vllm/platforms/cpu.py @bigPYJ1151 @xuechendi /cmake/cpu_extension.cmake @bigPYJ1151 -/docker/Dockerfile.cpu @bigPYJ1151 +/docker/Dockerfile.cpu @bigPYJ1151 @xuechendi # Intel GPU -/vllm/v1/worker/xpu* @jikunshang -/vllm/platforms/xpu.py @jikunshang -/docker/Dockerfile.xpu @jikunshang +/vllm/v1/worker/xpu* @jikunshang @xuechendi +/vllm/platforms/xpu.py @jikunshang @xuechendi +/docker/Dockerfile.xpu @jikunshang @xuechendi # Nemotron-specific files /vllm/model_executor/models/*nemotron* @tomeras91 diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 8043df65d558..8a3934670e44 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -15,7 +15,6 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT - [ ] The test plan, such as providing test command. - [ ] The test results, such as pasting the results comparison before and after, or e2e results - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. -- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0). **BEFORE SUBMITTING, PLEASE READ ** (anything written below this line will be removed by GitHub Actions) diff --git a/.github/mergify.yml b/.github/mergify.yml index e8ef4d49dd00..b96d6b81ac07 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -18,7 +18,7 @@ pull_request_rules: - name: comment-pre-commit-failure description: Comment on PR when pre-commit check fails conditions: - - status-failure=pre-commit + - check-failure=pre-commit - -closed - -draft actions: @@ -51,7 +51,7 @@ pull_request_rules: - name: comment-dco-failure description: Comment on PR when DCO check fails conditions: - - status-failure=dco + - check-failure=dco - -closed - -draft actions: @@ -83,8 +83,8 @@ pull_request_rules: - or: - files~=^examples/.*deepseek.*\.py - files~=^tests/.*deepseek.*\.py - - files~=^vllm/entrypoints/openai/tool_parsers/.*deepseek.*\.py - files~=^vllm/model_executor/models/.*deepseek.*\.py + - files~=^vllm/tool_parsers/.*deepseek.*\.py - files~=^vllm/reasoning/.*deepseek.*\.py - files~=^vllm/transformers_utils/.*deepseek.*\.py - title~=(?i)DeepSeek @@ -110,9 +110,10 @@ pull_request_rules: - or: - files~=^examples/.*llama.*\.py - files~=^tests/.*llama.*\.py - - files~=^vllm/entrypoints/openai/tool_parsers/llama.*\.py - files~=^vllm/model_executor/models/.*llama.*\.py - - files~=^vllm/transformers_utils/configs/.*llama.*\.py + - files~=^vllm/reasoning/.*llama.*\.py + - files~=^vllm/tool_parsers/.*llama.*\.py + - files~=^vllm/transformers_utils/.*llama.*\.py - title~=(?i)llama actions: label: @@ -133,6 +134,23 @@ pull_request_rules: add: - multi-modality +- name: label-mistral + description: Automatically apply mistral label + conditions: + - label != stale + - or: + - files~=^examples/.*mistral.*\.py + - files~=^tests/.*mistral.*\.py + - files~=^vllm/model_executor/models/.*mistral.*\.py + - files~=^vllm/reasoning/.*mistral.*\.py + - files~=^vllm/tool_parsers/.*mistral.*\.py + - files~=^vllm/transformers_utils/.*mistral.*\.py + - title~=(?i)Mistral + actions: + label: + add: + - mistral + - name: label-new-model description: Automatically apply new-model label conditions: @@ -167,7 +185,9 @@ pull_request_rules: - files~=^examples/.*qwen.*\.py - files~=^tests/.*qwen.*\.py - files~=^vllm/model_executor/models/.*qwen.*\.py + - files~=^vllm/tool_parsers/.*qwen.*\.py - files~=^vllm/reasoning/.*qwen.*\.py + - files~=^vllm/transformers_utils/.*qwen.*\.py - title~=(?i)Qwen actions: label: @@ -242,8 +262,9 @@ pull_request_rules: - files~=^docker/Dockerfile.xpu - files~=^\\.buildkite/intel_jobs/ - files=\.buildkite/ci_config_intel.yaml - - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py + - files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py + - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py - files=vllm/distributed/device_communicators/xpu_communicator.py - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -251,6 +272,7 @@ pull_request_rules: - files=vllm/v1/worker/xpu_worker.py - files=vllm/v1/worker/xpu_model_runner.py - files=vllm/_xpu_ops.py + - files=vllm/kernels/xpu_ops.py - files~=^vllm/lora/ops/xpu_ops - files=vllm/lora/punica_wrapper/punica_xpu.py - files=vllm/platforms/xpu.py @@ -258,7 +280,6 @@ pull_request_rules: - title~=(?i)XPU - title~=(?i)Intel - title~=(?i)BMG - - title~=(?i)Arc actions: label: add: @@ -378,17 +399,18 @@ pull_request_rules: add: - tool-calling -- name: auto-rebase if approved, ready, and 40 commits behind main +- name: auto-rebase to keep merge candidate within 1 day behind main conditions: - base = main - label=ready - "#approved-reviews-by >= 1" - - "#commits-behind >= 40" + - "#commits-behind >= 50" + - "#check-failure = 0" - -closed - -draft - -conflict actions: - rebase: {} + update: {} - name: ping author on conflicts and add 'needs-rebase' label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 2cb5c176ae0a..3efa582f6702 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -320,20 +320,25 @@ jobs: script: | // Configuration: Map labels to GitHub users to CC // You can add multiple users per label, and multiple label configurations + // {users} will be replaced with @mentions const ccConfig = { rocm: { - users: ['hongxiayang', 'tjtanaa', 'vllmellm'], // Add more users as needed: ['user1', 'user2', 'user3'] - message: 'CC {users} for ROCm-related issue' // {users} will be replaced with @mentions + users: ['hongxiayang', 'tjtanaa', 'vllmellm'], + message: 'CC {users} for ROCm-related issue', + }, + mistral: { + users: ['patrickvonplaten', 'juliendenize', 'andylolu2'], + message: 'CC {users} for Mistral-related issue', }, // Add more label -> user mappings here // Example: // cuda: { // users: ['user1', 'user2'], - // message: 'CC {users} for CUDA-related issue' + // message: 'CC {users} for CUDA-related issue', // }, // performance: { // users: ['perfexpert'], - // message: 'CC {users} for performance issue' + // message: 'CC {users} for performance issue', // }, }; diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index 3c1a50bf8085..ea1c8b0feac3 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -32,7 +32,7 @@ jobs: - name: Install dependencies and build vLLM run: | - uv pip install -r requirements/cpu-build.txt --index-strategy unsafe-best-match + uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match uv pip install -e . --no-build-isolation env: @@ -45,6 +45,7 @@ jobs: - name: Smoke test vllm serve run: | # Start server in background + VLLM_CPU_KVCACHE_SPACE=1 \ vllm serve Qwen/Qwen3-0.6B \ --max-model-len=2K \ --load-format=dummy \ diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index ef5e30952c62..27100f9f4da0 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -62,14 +62,14 @@ jobs: const prAuthor = context.payload.pull_request.user.login; const { data: searchResults } = await github.rest.search.issuesAndPullRequests({ - q: `repo:${owner}/${repo} type:pr author:${prAuthor}`, + q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`, per_page: 1, }); - const authorPRCount = searchResults.total_count; - console.log(`Found ${authorPRCount} PRs by ${prAuthor}`); + const mergedPRCount = searchResults.total_count; + console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`); - if (authorPRCount === 1) { + if (mergedPRCount === 0) { console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`); await github.rest.issues.createComment({ owner, @@ -98,5 +98,5 @@ jobs: ].join('\n'), }); } else { - console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`); + console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`); } diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index bbc55e9532b9..8ab8d3e7035f 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -2,6 +2,7 @@ name: pre-commit on: pull_request: + types: [opened, synchronize, reopened, labeled] push: branches: [main] @@ -15,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 @@ -44,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 diff --git a/.github/workflows/scripts/build.sh b/.github/workflows/scripts/build.sh index c69ebbb42da5..5bb9f7196806 100644 --- a/.github/workflows/scripts/build.sh +++ b/.github/workflows/scripts/build.sh @@ -9,12 +9,12 @@ PATH=${cuda_home}/bin:$PATH LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH # Install requirements -$python_executable -m pip install -r requirements/build.txt -r requirements/cuda.txt +$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt # Limit the number of parallel jobs to avoid OOM export MAX_JOBS=1 # Make sure release wheels are built for the following architectures -export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX" +export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" bash tools/check_repo.sh diff --git a/.gitignore b/.gitignore index 7b822165d3eb..134bbc5cc893 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,9 @@ vllm/third_party/triton_kernels/* # FlashMLA interface copied from source vllm/third_party/flashmla/flash_mla_interface.py +# DeepGEMM vendored package built from source +vllm/third_party/deep_gemm/ + # triton jit .triton @@ -26,6 +29,7 @@ __pycache__/ # Distribution / packaging .Python build/ +!requirements/build/ cmake-build-*/ CMakeUserPresets.json develop-eggs/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f55df24bc6f3..33b1db69dec4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,15 +39,24 @@ repos: rev: 0.11.1 hooks: - id: pip-compile - args: [requirements/test.in, -c, requirements/common.txt, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu130, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"] - files: ^requirements/test\.(in|txt)$ + args: [ + requirements/test/cuda.in, + -c, requirements/cuda.txt, + -o, requirements/test/cuda.txt, + --index-strategy, unsafe-best-match, + --torch-backend, cu130, + --python-platform, x86_64-manylinux_2_28, + --python-version, "3.12", + ] + files: ^requirements/(common|cuda|test/cuda)\.(in|txt)$ - id: pip-compile alias: pip-compile-rocm name: pip-compile-rocm args: [ - requirements/rocm-test.in, -o, requirements/rocm-test.txt, - --index-strategy, unsafe-best-match, + requirements/test/rocm.in, -c, requirements/rocm.txt, + -o, requirements/test/rocm.txt, + --index-strategy, unsafe-best-match, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12", # Exclude torch and CUDA/NVIDIA packages @@ -59,30 +68,76 @@ repos: --no-emit-package, cuda-pathfinder, --no-emit-package, cuda-toolkit, --no-emit-package, cupy-cuda12x, + # nvidia packages (unsuffixed / unified naming) --no-emit-package, nvidia-cublas, --no-emit-package, nvidia-cuda-cupti, --no-emit-package, nvidia-cuda-nvrtc, --no-emit-package, nvidia-cuda-runtime, - --no-emit-package, nvidia-cudnn-cu13, + --no-emit-package, nvidia-cudnn, --no-emit-package, nvidia-cufft, --no-emit-package, nvidia-cufile, --no-emit-package, nvidia-curand, --no-emit-package, nvidia-cusolver, --no-emit-package, nvidia-cusparse, + --no-emit-package, nvidia-cusparselt, + --no-emit-package, nvidia-nccl, + --no-emit-package, nvidia-nvjitlink, + --no-emit-package, nvidia-nvshmem, + --no-emit-package, nvidia-nvtx, + # nvidia cu12 packages + --no-emit-package, nvidia-cublas-cu12, + --no-emit-package, nvidia-cuda-cupti-cu12, + --no-emit-package, nvidia-cuda-nvrtc-cu12, + --no-emit-package, nvidia-cuda-runtime-cu12, + --no-emit-package, nvidia-cudnn-cu12, + --no-emit-package, nvidia-cufft-cu12, + --no-emit-package, nvidia-cufile-cu12, + --no-emit-package, nvidia-curand-cu12, + --no-emit-package, nvidia-cusolver-cu12, + --no-emit-package, nvidia-cusparse-cu12, + --no-emit-package, nvidia-cusparselt-cu12, + --no-emit-package, nvidia-nccl-cu12, + --no-emit-package, nvidia-nvjitlink-cu12, + --no-emit-package, nvidia-nvshmem-cu12, + --no-emit-package, nvidia-nvtx-cu12, + # nvidia cu13 packages + --no-emit-package, nvidia-cublas-cu13, + --no-emit-package, nvidia-cuda-cupti-cu13, + --no-emit-package, nvidia-cuda-nvrtc-cu13, + --no-emit-package, nvidia-cuda-runtime-cu13, + --no-emit-package, nvidia-cudnn-cu13, + --no-emit-package, nvidia-cufft-cu13, + --no-emit-package, nvidia-cufile-cu13, + --no-emit-package, nvidia-curand-cu13, + --no-emit-package, nvidia-cusolver-cu13, + --no-emit-package, nvidia-cusparse-cu13, --no-emit-package, nvidia-cusparselt-cu13, --no-emit-package, nvidia-nccl-cu13, - --no-emit-package, nvidia-nvjitlink, + --no-emit-package, nvidia-nvjitlink-cu13, --no-emit-package, nvidia-nvshmem-cu13, - --no-emit-package, nvidia-nvtx, + --no-emit-package, nvidia-nvtx-cu13, + ] + files: ^requirements/(common|rocm|test/rocm)\.(in|txt)$ + - id: pip-compile + alias: pip-compile-xpu + name: pip-compile-xpu + args: [ + requirements/test/xpu.in, + -c, requirements/xpu.txt, + -o, requirements/test/xpu.txt, + --index-strategy, unsafe-best-match, + --torch-backend, xpu, + --python-platform, x86_64-manylinux_2_39, + --python-version, "3.12", ] - files: ^requirements/rocm-test\.(in|txt)$ + files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$ - repo: local hooks: - id: format-torch-nightly-test - name: reformat nightly_torch_test.txt to be in sync with test.in + name: reformat test/nightly-torch.txt to be in sync with test/cuda.in language: python entry: python tools/pre_commit/generate_nightly_torch_test.py - files: ^requirements/test\.(in|txt)$ + files: ^requirements/test/cuda\.(in|txt)$ - id: mypy-local name: Run mypy locally for lowest supported Python version entry: python tools/pre_commit/mypy.py 0 "3.10" diff --git a/AGENTS.md b/AGENTS.md index 61312b29ef7d..215e4195eb43 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -72,11 +72,11 @@ uv pip install -e . --torch-backend=auto ```bash # Install test dependencies. -# requirements/test.txt is pinned to x86_64; on other platforms, use the +# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the # unpinned source file instead: -uv pip install -r requirements/test.in # resolves for current platform +uv pip install -r requirements/test/cuda.in # resolves for current platform # Or on x86_64: -uv pip install -r requirements/test.txt +uv pip install -r requirements/test/cuda.txt # Run a specific test file (use .venv/bin/python directly; # `source activate` does not persist in non-interactive shells): diff --git a/CMakeLists.txt b/CMakeLists.txt index 48bd2294cef4..fb8a1d7e1e14 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) # Supported python versions. These versions will be searched in order, the # first match will be selected. These should be kept in sync with setup.py. # -set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13") +set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. @@ -94,12 +94,15 @@ find_package(Torch REQUIRED) # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) - set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1") + # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc) + # to support the whole generation without specifying all sub-architectures + # see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/ + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0") elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8) - set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1") + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1") else() - set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0") + set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0") endif() # @@ -225,8 +228,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") # Certain HIP functions are marked as [[nodiscard]], yet vllm ignores the result which generates # a lot of warnings that always mask real issues. Suppressing until this is properly addressed. # - set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result") + set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result -Wno-unused-value") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result -Wno-unused-value") endif() # @@ -299,6 +302,7 @@ set(VLLM_EXT_SRC "csrc/quantization/w8a8/int8/scaled_quant.cu" "csrc/quantization/w8a8/fp8/common.cu" "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" "csrc/quantization/gguf/gguf_kernel.cu" "csrc/quantization/activation_kernels.cu" "csrc/cuda_utils_kernels.cu" @@ -306,6 +310,10 @@ set(VLLM_EXT_SRC "csrc/torch_bindings.cpp") if(VLLM_GPU_LANG STREQUAL "CUDA") + list(APPEND VLLM_EXT_SRC + "csrc/minimax_reduce_rms_kernel.cu" + "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. @@ -340,8 +348,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_EXT_SRC "csrc/quantization/awq/gemm_kernels.cu" - "csrc/cutlass_extensions/common.cpp" - "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu") + "csrc/cutlass_extensions/common.cpp") set_gencode_flags_for_srcs( SRCS "${VLLM_EXT_SRC}" @@ -921,6 +928,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly + # from cache_kernels.cu, so it belongs in _C rather than _C_stable. + set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set_gencode_flags_for_srcs( + SRCS "${NVFP4_KV_SRC}" + CUDA_ARCHS "${FP4_ARCHS}") + target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}") @@ -942,11 +957,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu") + "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu") set_gencode_flags_for_srcs( SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set_gencode_flags_for_srcs( + SRCS "${NVFP4_KV_SRC}" + CUDA_ARCHS "${FP4_ARCHS}") + target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}") @@ -1030,7 +1053,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu" "csrc/moe/grouped_topk_kernels.cu" - "csrc/moe/router_gemm.cu") + "csrc/moe/router_gemm.cu" + "csrc/moe/topk_softplus_sqrt_kernels.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") @@ -1222,6 +1246,7 @@ endif() # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") + include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/flashmla.cmake) include(cmake/external_projects/qutlass.cmake) diff --git a/README.md b/README.md index 44418b712a26..d5653c76c8d5 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone! | Documentation | Blog | Paper | Twitter/X | User Forum | Developer Slack |

-🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more. +🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more. For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us. --- @@ -23,47 +23,54 @@ For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us. vLLM is a fast and easy-to-use library for LLM inference and serving. -Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has evolved into a community-driven project with contributions from both academia and industry. +Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors. vLLM is fast with: - State-of-the-art serving throughput - Efficient management of attention key and value memory with [**PagedAttention**](https://blog.vllm.ai/2023/06/20/vllm.html) -- Continuous batching of incoming requests -- Fast model execution with CUDA/HIP graph -- Quantizations: [GPTQ](https://arxiv.org/abs/2210.17323), [AWQ](https://arxiv.org/abs/2306.00978), [AutoRound](https://arxiv.org/abs/2309.05516), INT4, INT8, and FP8 -- Optimized CUDA kernels, including integration with FlashAttention and FlashInfer -- Speculative decoding -- Chunked prefill +- Continuous batching of incoming requests, chunked prefill, prefix caching +- Fast and flexible model execution with piecewise and full CUDA/HIP graphs +- Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and [more](https://docs.vllm.ai/en/latest/features/quantization/index.html) +- Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton +- Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL +- Speculative decoding including n-gram, suffix, EAGLE, DFlash +- Automatic kernel generation and graph-level transformations using torch.compile +- Disaggregated prefill, decode, and encode vLLM is flexible and easy to use with: - Seamless integration with popular Hugging Face models - High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more -- Tensor, pipeline, data and expert parallelism support for distributed inference +- Tensor, pipeline, data, expert, and context parallelism for distributed inference - Streaming outputs -- OpenAI-compatible API server -- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, Arm CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend. -- Prefix caching support -- Multi-LoRA support +- Generation of structured outputs using xgrammar or guidance +- Tool calling and reasoning parsers +- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support +- Efficient multi-LoRA support for dense and MoE layers +- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. -vLLM seamlessly supports most popular open-source models on HuggingFace, including: +vLLM seamlessly supports 200+ model architectures on Hugging Face, including: -- Transformer-like LLMs (e.g., Llama) -- Mixture-of-Expert LLMs (e.g., Mixtral, Deepseek-V2 and V3) -- Embedding Models (e.g., E5-Mistral) -- Multi-modal LLMs (e.g., LLaVA) +- Decoder-only LLMs (e.g., Llama, Qwen, Gemma) +- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS) +- Hybrid attention and state-space models (e.g., Mamba, Qwen3.5) +- Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral) +- Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT) +- Reward and classification models (e.g., Qwen-Math) Find the full list of supported models [here](https://docs.vllm.ai/en/latest/models/supported_models.html). ## Getting Started -Install vLLM with `pip` or [from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source): +Install vLLM with [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip`: ```bash -pip install vllm +uv pip install vllm ``` +Or [build from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for development. + Visit our [documentation](https://docs.vllm.ai/en/latest/) to learn more. - [Installation](https://docs.vllm.ai/en/latest/getting_started/installation.html) diff --git a/tests/entrypoints/pooling/pooling/__init__.py b/benchmarks/__init__.py similarity index 100% rename from tests/entrypoints/pooling/pooling/__init__.py rename to benchmarks/__init__.py diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index f8bc7b4a10ed..b58ddffcd0b1 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -404,6 +404,7 @@ def _build_attention_metadata( query_start_loc=q_start_gpu, query_start_loc_cpu=q_start_cpu, seq_lens=seq_lens_gpu, + seq_lens_cpu_upper_bound=seq_lens_cpu, _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, slot_mapping=slot_mapping, diff --git a/benchmarks/kernels/__init__.py b/benchmarks/kernels/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/benchmarks/kernels/benchmark_block_fp8_gemm.py b/benchmarks/kernels/benchmark_block_fp8_gemm.py index 8d50c3828206..9eddc907b937 100644 --- a/benchmarks/kernels/benchmark_block_fp8_gemm.py +++ b/benchmarks/kernels/benchmark_block_fp8_gemm.py @@ -9,11 +9,12 @@ import torch from vllm.benchmarks.lib.utils import default_vllm_config -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - W8A8BlockFp8LinearOp, +from vllm.model_executor.kernels.linear import ( + init_fp8_linear_kernel, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + create_fp8_quant_key, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( CUTLASS_BLOCK_FP8_SUPPORTED, @@ -70,11 +71,15 @@ def build_w8a8_block_fp8_runner(M, N, K, block_size, device, use_cutlass): weight_group_shape = GroupShape(block_n, block_k) act_quant_group_shape = GroupShape(1, block_k) # Per-token, per-group quantization - linear_op = W8A8BlockFp8LinearOp( - weight_group_shape=weight_group_shape, - act_quant_group_shape=act_quant_group_shape, - cutlass_block_fp8_supported=use_cutlass, - use_aiter_and_is_supported=False, + linear_op = init_fp8_linear_kernel( + weight_quant_key=create_fp8_quant_key( + static=True, group_shape=weight_group_shape + ), + activation_quant_key=create_fp8_quant_key( + static=False, group_shape=act_quant_group_shape + ), + out_dtype=torch.get_default_dtype(), + module_name="build_w8a8_block_fp8_runner", ) def run(): diff --git a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py index 3f80b024e108..03d7fb386f74 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_fp8.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_fp8.py @@ -16,7 +16,7 @@ maybe_make_prepare_finalize, ) from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config -from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 +from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser diff --git a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py index 2d4afd38c097..7379bf858889 100644 --- a/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py +++ b/benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py @@ -22,7 +22,7 @@ fp8_w8a8_moe_quant_config, nvfp4_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.cutlass_moe import ( +from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp4, ) from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk diff --git a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py index dd4060bbdb94..04fc2960d1e4 100644 --- a/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py +++ b/benchmarks/kernels/benchmark_grouped_gemm_cutlass.py @@ -13,7 +13,7 @@ maybe_make_prepare_finalize, ) from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config -from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8 +from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8 from vllm.model_executor.layers.fused_moe.fused_moe import ( fused_experts, fused_topk, diff --git a/benchmarks/kernels/benchmark_moe_align_block_size.py b/benchmarks/kernels/benchmark_moe_align_block_size.py index 5f9a131f79b0..a340500379a0 100644 --- a/benchmarks/kernels/benchmark_moe_align_block_size.py +++ b/benchmarks/kernels/benchmark_moe_align_block_size.py @@ -9,6 +9,7 @@ moe_align_block_size, ) from vllm.triton_utils import triton +from vllm.utils.torch_utils import set_random_seed def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor: @@ -44,7 +45,7 @@ def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor: def benchmark(num_tokens, num_experts, topk, ep_size, provider): """Benchmark function for Triton.""" block_size = 256 - torch.cuda.manual_seed_all(0) + set_random_seed(0) topk_ids = get_topk_ids(num_tokens, num_experts, topk) e_map = None diff --git a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py index 13b97b7696b3..9fcf278f2ef3 100644 --- a/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py +++ b/benchmarks/kernels/benchmark_silu_mul_fp8_quant.py @@ -20,7 +20,7 @@ import numpy as np import torch -from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import ( +from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( persistent_masked_m_silu_mul_quant, ) from vllm.triton_utils import tl, triton diff --git a/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py b/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py index 4384d3e56828..c9aaef284d70 100644 --- a/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py +++ b/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py @@ -16,6 +16,7 @@ fp8_gemm_nt, per_block_cast_to_fp8, ) +from vllm.utils.torch_utils import set_random_seed def benchmark_shape( @@ -235,9 +236,7 @@ def run_benchmarks(verbose: bool = False): torch.backends.cudnn.allow_tf32 = True # Set seeds for reproducibility - torch.manual_seed(42) - torch.cuda.manual_seed(42) - + set_random_seed(42) # Define benchmark shapes (m, n, k) shapes = [ (8, 4096, 7168), diff --git a/benchmarks/kernels/ir/__init__.py b/benchmarks/kernels/ir/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/benchmarks/kernels/ir/bench_ir_ops.py b/benchmarks/kernels/ir/bench_ir_ops.py new file mode 100644 index 000000000000..b23c4e8ae327 --- /dev/null +++ b/benchmarks/kernels/ir/bench_ir_ops.py @@ -0,0 +1,378 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Generic benchmark harness for vLLM IR ops. + +Usage: + python benchmarks/kernels/ir/bench_ir_ops.py + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm,silu_mul + python benchmarks/kernels/ir/bench_ir_ops.py --no-cuda-graph + python benchmarks/kernels/ir/bench_ir_ops.py --ops rms_norm --save-path ./results/ +""" + +import argparse +import contextlib +import csv +import dataclasses +import datetime +import math +import os +import subprocess +import sys +import tempfile + +# Ensure repo root is on sys.path so `benchmarks` is importable as a package. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +# Suppress noisy C++ warnings from vllm kernel registration (written to fd 2 +# directly by the dynamic linker, so Python-level sys.stderr redirect won't +# catch them). +_saved_fd = os.dup(2) +try: + with open(os.devnull, "w") as _devnull: + os.dup2(_devnull.fileno(), 2) + import torch + + import vllm.kernels # noqa: E402, F401 +finally: + os.dup2(_saved_fd, 2) + os.close(_saved_fd) + +from tqdm import tqdm # noqa: E402 + +from benchmarks.kernels.ir.shapes import SHAPE_CONFIGS # noqa: E402 # isort: skip +from vllm.ir.op import IrOp # noqa: E402 +from vllm.platforms import current_platform # noqa: E402 +from vllm.triton_utils import triton # noqa: E402 + + +@dataclasses.dataclass(frozen=True) +class BenchConfig: + use_cuda_graph: bool = True + warmup: int = 25 + rep: int = 100 + + +def _pkg_version(name: str) -> str: + from importlib.metadata import PackageNotFoundError, version + + with contextlib.suppress(PackageNotFoundError): + return version(name) + return "not installed" + + +_METADATA_LABELS = { + "timestamp": "Timestamp", + "git_commit": "Git commit", + "vllm": "vLLM", + "pytorch": "PyTorch", + "cuda_runtime": "CUDA runtime", + "triton": "Triton", + "cutlass": "CUTLASS", + "helion": "Helion", + "device": "Device", + "bench_mode": "Bench mode", + "warmup": "Warmup", + "rep": "Repetitions", +} + + +def collect_env_metadata(cfg: BenchConfig) -> dict[str, str]: + from vllm.collect_env import get_env_info + + env = get_env_info() + + git_sha = "unknown" + with contextlib.suppress(subprocess.CalledProcessError, FileNotFoundError): + git_sha = ( + subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL + ) + .decode() + .strip() + ) + + device_name = current_platform.get_device_name() + + warmup_note = " ms" if not cfg.use_cuda_graph else " ms (ignored)" + rep_note = " replays" if cfg.use_cuda_graph else " ms" + + return { + "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + "git_commit": git_sha, + "vllm": str(env.vllm_version), + "pytorch": str(env.torch_version), + "cuda_runtime": str(env.cuda_runtime_version), + "triton": triton.__version__, + "cutlass": _pkg_version("nvidia-cutlass-dsl"), + "helion": _pkg_version("helion"), + "device": device_name, + "bench_mode": "cuda_graph" if cfg.use_cuda_graph else "eager", + "warmup": f"{cfg.warmup}{warmup_note}", + "rep": f"{cfg.rep}{rep_note}", + } + + +def print_metadata(metadata: dict[str, str]): + print("=" * 60) + for key, val in metadata.items(): + print(f"{_METADATA_LABELS.get(key, key) + ':':<16}{val}") + print("=" * 60) + + +def _clone_args(args: tuple) -> tuple: + return tuple(a.clone() if isinstance(a, torch.Tensor) else a for a in args) + + +# TODO(gmagogsfm): When the `maybe_inplace` PR lands, ops marked as +# inplace=True will mutate bench_args across iterations. Both CUDA graph +# and eager modes will accumulate drift from repeated in-place mutation. +# We need to re-clone inputs per iteration for inplace ops. +def _bench_one(fn, args, cfg: BenchConfig) -> float: + bench_args = _clone_args(args) + bench_fn = lambda: fn(*bench_args) + + if cfg.use_cuda_graph: + ms = triton.testing.do_bench_cudagraph(bench_fn, rep=cfg.rep, quantiles=[0.5]) + else: + ms = triton.testing.do_bench( + bench_fn, warmup=cfg.warmup, rep=cfg.rep, quantiles=[0.5] + ) + return ms * 1000 + + +# TODO(gmagogsfm): Once compiled native implementation lands (#38775), +# the benchmark baseline should be the compiled native (what vLLM runs by +# default) rather than the uncompiled native implementation. +def collect_timings( + op: IrOp, shape_configs: list[dict], cfg: BenchConfig +) -> tuple[list[str], list[str], dict[str, dict[str, float]]]: + def fmt(v) -> str: + return str(v).split(".")[-1] if isinstance(v, torch.dtype) else str(v) + + case_names = [ + "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs + ] + providers = [n for n, impl in op.impls.items() if impl.supported] + + results: dict[str, dict[str, float]] = {c: {} for c in case_names} + for provider in providers: + impl = op.impls[provider] + desc = f"{op.name} / {provider}" + for case_name, kwargs in tqdm( + zip(case_names, shape_configs), + desc=desc, + total=len(case_names), + unit=" cases", + ): + args = op.generate_inputs(**kwargs) + if impl.supports_args(*args): + results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg) + else: + results[case_name][provider] = float("nan") + + return case_names, providers, results + + +def analyze_results( + op_name: str, + case_names: list[str], + providers: list[str], + results: dict[str, dict[str, float]], +) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]: + native_col = "native" + non_native = [p for p in providers if p != native_col] + + header_cols = ["case"] + for p in providers: + header_cols.append(f"{p} (us)") + for p in non_native: + header_cols.append(f"{p} speedup") + + detail_rows: list[dict[str, str]] = [] + speedup_data: dict[str, list[tuple[float, str]]] = {p: [] for p in non_native} + + for case_name in case_names: + timings = results[case_name] + row: dict[str, str] = {"case": case_name} + + for p in providers: + val = timings.get(p, float("nan")) + row[f"{p} (us)"] = f"{val:.2f}" if not math.isnan(val) else "n/a" + + native_us = timings.get(native_col, float("nan")) + for p in non_native: + p_us = timings.get(p, float("nan")) + if not math.isnan(native_us) and not math.isnan(p_us) and p_us > 0: + speedup = native_us / p_us + row[f"{p} speedup"] = f"{speedup:.2f}x" + speedup_data[p].append((speedup, case_name)) + else: + row[f"{p} speedup"] = "n/a" + + detail_rows.append(row) + + summary_rows: list[dict[str, str]] = [] + for p in non_native: + entries = speedup_data[p] + if not entries: + continue + speedups = [s for s, _ in entries] + geomean = math.exp(sum(math.log(s) for s in speedups) / len(speedups)) + best_val, best_case = max(entries) + worst_val, worst_case = min(entries) + wins = sum(1 for s in speedups if s > 1.0) + losses = sum(1 for s in speedups if s < 1.0) + total = len(speedups) + + print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):") + print(f" geomean speedup: {geomean:.2f}x") + print(f" best: {best_val:.2f}x ({best_case})") + print(f" worst: {worst_val:.2f}x ({worst_case})") + + summary_rows.append( + { + "op": op_name, + "provider": p, + "geomean_speedup": f"{geomean:.2f}", + "best_speedup": f"{best_val:.2f}", + "best_case": best_case, + "worst_speedup": f"{worst_val:.2f}", + "worst_case": worst_case, + "wins": str(wins), + "losses": str(losses), + "total": str(total), + } + ) + + return detail_rows, summary_rows, header_cols + + +def write_csv(path: str, rows: list[dict[str, str]], fieldnames: list[str]): + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def save_results( + save_dir: str, + op_name: str, + detail_rows: list[dict[str, str]], + header_cols: list[str], + all_summary_rows: list[dict[str, str]], + metadata: dict[str, str], +): + write_csv( + os.path.join(save_dir, f"{op_name}_detail.csv"), + detail_rows, + header_cols, + ) + if all_summary_rows: + write_csv( + os.path.join(save_dir, "summary.csv"), + all_summary_rows, + list(all_summary_rows[0].keys()), + ) + write_csv( + os.path.join(save_dir, "metadata.csv"), + [metadata], + list(metadata.keys()), + ) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops") + parser.add_argument( + "--ops", + type=str, + default=None, + help="Comma-separated list of op names to benchmark (substring match)", + ) + parser.add_argument( + "--no-cuda-graph", + action="store_true", + help="Disable CUDA graph; use do_bench with L2 cache flushing instead", + ) + parser.add_argument( + "--warmup", + type=int, + default=25, + help="Warmup time in ms (do_bench) or ignored with CUDA graph (default: 25)", + ) + parser.add_argument( + "--rep", + type=int, + default=100, + help="Repetition time in ms (do_bench) or number of graph replays " + "(do_bench_cudagraph) (default: 100)", + ) + parser.add_argument( + "--save-path", + type=str, + default=None, + help="Directory to save results (default: auto-created temp dir)", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + cfg = BenchConfig( + use_cuda_graph=not args.no_cuda_graph, + warmup=args.warmup, + rep=args.rep, + ) + + torch.set_default_device(current_platform.device_type) + + metadata = collect_env_metadata(cfg) + print_metadata(metadata) + + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + save_dir = args.save_path or os.path.join( + tempfile.gettempdir(), f"vllm_ir_bench_{timestamp}" + ) + os.makedirs(save_dir, exist_ok=True) + + op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None + all_summary_rows: list[dict[str, str]] = [] + + for op in IrOp.registry.values(): + if op_filters and not any(f in op.name for f in op_filters): + continue + if not op.has_input_generator: + print(f"Skipping op '{op.name}': no input generator registered") + continue + if op.name not in SHAPE_CONFIGS: + raise RuntimeError( + f"No benchmark shape config for op '{op.name}'. " + f"Add it to benchmarks/kernels/ir/shapes.py" + ) + + case_names, providers, results = collect_timings( + op, SHAPE_CONFIGS[op.name], cfg + ) + detail_rows, summary_rows, header_cols = analyze_results( + op.name, case_names, providers, results + ) + all_summary_rows.extend(summary_rows) + + save_results( + save_dir, + op.name, + detail_rows, + header_cols, + all_summary_rows, + metadata, + ) + + print(f"\nResults saved to: {save_dir}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/ir/shapes.py b/benchmarks/kernels/ir/shapes.py new file mode 100644 index 000000000000..6cc44cf6cec1 --- /dev/null +++ b/benchmarks/kernels/ir/shapes.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shape configurations for IR op benchmarks. +""" + +import torch + +NUM_TOKENS = [1, 2, 4, 16, 64, 256, 1024, 4096, 16384] +COMMON_HIDDEN_SIZES = [ + 2048, # Llama 3.2 1B, Qwen 3 MoE 30B-A3B, Gemma 3n + 3072, # Gemma 7B/9B + 4096, # Llama 3 8B, Qwen 3 8B, Mistral 7B + 5120, # Llama 4 Scout 17B-16E + 7168, # DeepSeek V3 + 8192, # Llama 3 70B + 16384, # Llama 3 405B +] + +# Each entry maps an op name to a list of kwarg dicts that will be passed +# to that op's registered input generator via op.generate_inputs(**kwargs). +SHAPE_CONFIGS: dict[str, list[dict]] = { + "rms_norm": [ + {"num_tokens": n, "hidden_size": d, "dtype": dtype} + for dtype in [torch.float16, torch.bfloat16, torch.float32] + for d in COMMON_HIDDEN_SIZES + for n in NUM_TOKENS + ], +} diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index e23f6b923f1b..881039f43f07 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -1439,6 +1439,12 @@ async def main() -> None: action="store_true", help="Export summary to Excel file (optional)", ) + parser.add_argument( + "--stats-json-output", + type=str, + default=None, + help="Export per-request stats (ttft_ms, tpot_ms, etc.) to a JSON file", + ) parser.add_argument( "-v", "--verbose", @@ -1651,6 +1657,19 @@ async def main() -> None: warmup_runtime_sec=warmup_runtime_sec, ) + if args.stats_json_output is not None: + # Export per-request metrics as a JSON array for downstream analysis. + stats_data = [s._asdict() for s in client_metrics] + logger.info( + f"{Color.GREEN}Writing per-request stats JSON: " + f"{args.stats_json_output}{Color.RESET}" + ) + os.makedirs( + os.path.dirname(os.path.abspath(args.stats_json_output)), exist_ok=True + ) + with open(args.stats_json_output, "w") as f: + json.dump(stats_data, f, indent=2) + if args.output_file is not None: # Write a JSON file with the updated conversations # The "assistant" content will contain the answers from the tested LLM diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 1b3f0d5ad833..8535186cc1ec 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -30,6 +30,21 @@ else() list(APPEND CXX_COMPILE_FLAGS "-fopenmp" "-DVLLM_CPU_EXTENSION") + + # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0) + # and create a local shim dir with it + vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR) + + find_library(OPEN_MP + NAMES gomp + PATHS ${VLLM_TORCH_GOMP_SHIM_DIR} + NO_DEFAULT_PATH + REQUIRED + ) + # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch + if (OPEN_MP) + set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}") + endif() endif() if (NOT MACOSX_FOUND) @@ -146,16 +161,49 @@ elseif (S390_FOUND) "-mtune=native") elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "RISC-V detected") - if(RVV_BF16_FOUND) - message(STATUS "BF16 extension detected") - set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d) - add_compile_definitions(RISCV_BF16_SUPPORT) - elseif (RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") - set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d) + # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo + # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256. + if(NOT DEFINED VLLM_RVV_VLEN) + # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. + if(EXISTS /proc/cpuinfo) + file(READ /proc/cpuinfo _cpuinfo) + set(_best 0) + foreach(_n IN ITEMS 128 256 512 1024) + if(_cpuinfo MATCHES "zvl${_n}b") + set(_best ${_n}) + endif() + endforeach() + if(_best GREATER 0) + set(VLLM_RVV_VLEN ${_best}) + endif() + endif() + # If auto-detect failed (no /proc/cpuinfo or no zvlb reported) + # but the compiler supports RVV, require explicit specification. + if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) + message(FATAL_ERROR + "RISC-V RVV is available but VLEN could not be auto-detected. " + "Please specify VLEN explicitly:\n" + " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" + " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n" + " -DVLLM_RVV_VLEN=0 (force scalar, no RVV)") + endif() + endif() + if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) + message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}") + if(RVV_BF16_FOUND) + message(STATUS "BF16 extension detected") + set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) + add_compile_definitions(RISCV_BF16_SUPPORT) + elseif(RVV_FP16_FOUND) + message(WARNING "BF16 functionality is not available") + set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) + else() + message(STATUS "compile riscv with scalar (no FP16/BF16)") + set(MARCH_FLAGS -march=rv64gc) + endif() else() message(STATUS "compile riscv with scalar") - list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc") + set(MARCH_FLAGS -march=rv64gc) endif() list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS}) else() @@ -175,20 +223,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND if(NOT NPROC) set(NPROC 4) endif() - # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0) - # and create a local shim dir with it - vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR) - - find_library(OPEN_MP - NAMES gomp - PATHS ${VLLM_TORCH_GOMP_SHIM_DIR} - NO_DEFAULT_PATH - REQUIRED - ) - # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch - if (OPEN_MP) - set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}") - endif() # Fetch and populate ACL if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}") @@ -349,6 +383,7 @@ endif() set(VLLM_EXT_SRC "csrc/cpu/activation.cpp" "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" @@ -359,6 +394,7 @@ set(VLLM_EXT_SRC if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) set(VLLM_EXT_SRC "csrc/cpu/shm.cpp" + "csrc/cpu/activation_lut_bf16.cpp" ${VLLM_EXT_SRC}) endif() @@ -383,6 +419,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/cpu_wna16.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/torch_bindings.cpp" @@ -395,6 +432,7 @@ if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_AVX2 "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake new file mode 100644 index 000000000000..0d7ea43fb7d0 --- /dev/null +++ b/cmake/external_projects/deepgemm.cmake @@ -0,0 +1,156 @@ +include(FetchContent) + +# If DEEPGEMM_SRC_DIR is set, DeepGEMM is built from that directory +# instead of downloading. +# It can be set as an environment variable or passed as a cmake argument. +# The environment variable takes precedence. +if (DEFINED ENV{DEEPGEMM_SRC_DIR}) + set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR}) +endif() + +if(DEEPGEMM_SRC_DIR) + FetchContent_Declare( + deepgemm + SOURCE_DIR ${DEEPGEMM_SRC_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +else() + # This ref should be kept in sync with tools/install_deepgemm.sh + FetchContent_Declare( + deepgemm + GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git + GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f + GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_PROGRESS TRUE + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +endif() + +# Use FetchContent_Populate (not MakeAvailable) to avoid processing +# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls. +FetchContent_GetProperties(deepgemm) +if(NOT deepgemm_POPULATED) + FetchContent_Populate(deepgemm) +endif() +message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}") + +# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 +set(DEEPGEMM_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a") +endif() +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") +elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") +endif() + +cuda_archs_loose_intersection(DEEPGEMM_ARCHS + "${DEEPGEMM_SUPPORT_ARCHS}" "${CUDA_ARCHS}") + +if(DEEPGEMM_ARCHS) + message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}") + + find_package(CUDAToolkit REQUIRED) + + # + # Build the _C pybind11 extension from DeepGEMM's C++ source. + # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime. + # + Python_add_library(_deep_gemm_C MODULE WITH_SOABI + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") + + # The pybind11 module name must be _C to match DeepGEMM's Python imports. + set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C") + + target_compile_definitions(_deep_gemm_C PRIVATE + "-DTORCH_EXTENSION_NAME=_C") + + target_include_directories(_deep_gemm_C PRIVATE + "${deepgemm_SOURCE_DIR}/csrc" + "${deepgemm_SOURCE_DIR}/deep_gemm/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include" + "${deepgemm_SOURCE_DIR}/third-party/fmt/include") + + target_compile_options(_deep_gemm_C PRIVATE + $<$:-std=c++17> + $<$:-O3> + $<$:-Wno-psabi> + $<$:-Wno-deprecated-declarations>) + + # torch_python is required because DeepGEMM uses pybind11 type casters + # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which + # use torch::Library custom ops. + find_library(TORCH_PYTHON_LIBRARY torch_python + PATHS "${TORCH_INSTALL_PREFIX}/lib" + REQUIRED) + + target_link_libraries(_deep_gemm_C PRIVATE + torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}" + CUDA::cudart CUDA::nvrtc) + + # Install the shared library into the vendored package directory + install(TARGETS _deep_gemm_C + LIBRARY DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C) + + # + # Vendor DeepGEMM Python package files + # + install(FILES + "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py" + DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C) + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/utils/" + DESTINATION vllm/third_party/deep_gemm/utils + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/testing/" + DESTINATION vllm/third_party/deep_gemm/testing + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/legacy/" + DESTINATION vllm/third_party/deep_gemm/legacy + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/" + DESTINATION vllm/third_party/deep_gemm/mega + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + # Generate envs.py (normally generated by DeepGEMM's setup.py build step) + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py" + "# Pre-installed environment variables\npersistent_envs = dict()\n") + install(FILES "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py" + DESTINATION vllm/third_party/deep_gemm + RENAME envs.py + COMPONENT _deep_gemm_C) + + # + # Install include files needed for JIT compilation at runtime. + # The JIT compiler finds these relative to the package directory. + # + + # DeepGEMM's own CUDA headers + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/include/" + DESTINATION vllm/third_party/deep_gemm/include + COMPONENT _deep_gemm_C) + + # CUTLASS and CuTe headers (vendored for JIT, separate from vLLM's CUTLASS) + install(DIRECTORY "${deepgemm_SOURCE_DIR}/third-party/cutlass/include/" + DESTINATION vllm/third_party/deep_gemm/include + COMPONENT _deep_gemm_C) + +else() + message(STATUS "DeepGEMM will not compile: " + "unsupported CUDA architecture ${CUDA_ARCHS}") + # Create empty target so setup.py doesn't fail on unsupported systems + add_custom_target(_deep_gemm_C) +endif() diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 0f16b9161fa3..65986df55012 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1 + GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 9414f5af3722..4aef7282134b 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG c0ec424fd8a546d0cbbf4bf050bbcfe837c55afb + GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn @@ -87,18 +87,30 @@ endforeach() # add_custom_target(_vllm_fa4_cutedsl_C) -# Copy flash_attn/cute directory (needed for FA4) and transform imports -# The cute directory uses flash_attn.cute imports internally, which we replace -# with vllm.vllm_flash_attn.cute to match our package structure. -install(CODE " - file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\") - foreach(SRC_FILE \${CUTE_PY_FILES}) - file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE}) - set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\") - get_filename_component(DST_DIR \${DST_FILE} DIRECTORY) - file(MAKE_DIRECTORY \${DST_DIR}) - file(READ \${SRC_FILE} FILE_CONTENTS) - string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\") - file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\") - endforeach() -" COMPONENT _vllm_fa4_cutedsl_C) +# Install flash_attn/cute directory (needed for FA4). +# When using a local source dir (VLLM_FLASH_ATTN_SRC_DIR), create a symlink +# so edits to cute-dsl Python files take effect immediately without rebuilding. +# Otherwise, copy files and transform flash_attn.cute imports to +# vllm.vllm_flash_attn.cute to match our package structure. +if(VLLM_FLASH_ATTN_SRC_DIR) + install(CODE " + set(LINK_TARGET \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\") + set(LINK_NAME \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute\") + file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\") + file(REMOVE_RECURSE \"\${LINK_NAME}\") + file(CREATE_LINK \"\${LINK_TARGET}\" \"\${LINK_NAME}\" SYMBOLIC) + " COMPONENT _vllm_fa4_cutedsl_C) +else() + install(CODE " + file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\") + foreach(SRC_FILE \${CUTE_PY_FILES}) + file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE}) + set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\") + get_filename_component(DST_DIR \${DST_FILE} DIRECTORY) + file(MAKE_DIRECTORY \${DST_DIR}) + file(READ \${SRC_FILE} FILE_CONTENTS) + string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\") + file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\") + endforeach() + " COMPONENT _vllm_fa4_cutedsl_C) +endif() diff --git a/csrc/activation_kernels.cu b/csrc/activation_kernels.cu index 758a77795553..303433392c32 100644 --- a/csrc/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -11,29 +11,74 @@ namespace vllm { template + bool act_first, bool HAS_CLAMP> __device__ __forceinline__ scalar_t compute(const scalar_t& x, - const scalar_t& y) { - return act_first ? ACT_FN(x) * y : x * ACT_FN(y); + 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); + } } template + bool act_first, bool HAS_CLAMP> __device__ __forceinline__ packed_t packed_compute(const packed_t& x, - const packed_t& y) { - return act_first ? packed_mul(PACKED_ACT_FN(x), y) - : packed_mul(x, PACKED_ACT_FN(y)); + 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(g); + up = cast_to_packed(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(g); + up = cast_to_packed(u); + } + return packed_mul(gate, PACKED_ACT_FN(up)); + } } // Activation and gating kernel template. template + bool use_vec, bool HAS_CLAMP, 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 int d, const float limit) { 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; @@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel( } #pragma unroll for (int j = 0; j < pvec_t::NUM_ELTS; j++) { - x.elts[j] = packed_compute( - x.elts[j], y.elts[j]); + x.elts[j] = + packed_compute( + x.elts[j], y.elts[j], limit); } if constexpr (use_256b) { st256(x, &out_vec[i]); @@ -72,7 +118,8 @@ __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(x, y); + out_ptr[idx] = + compute(x, y, limit); } } } @@ -151,8 +198,11 @@ 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. -#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \ +// 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) \ auto dtype = input.scalar_type(); \ int d = input.size(-1) / 2; \ int64_t num_tokens = input.numel() / input.size(-1); \ @@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) { scalar_t, typename vllm::PackedTypeConverter::Type, \ KERNEL, \ PACKED_KERNEL::Type>, \ - ACT_FIRST, true, true><<>>( \ - out.data_ptr(), input.data_ptr(), d); \ + ACT_FIRST, true, HAS_CLAMP, true><<>>( \ + out.data_ptr(), input.data_ptr(), d, LIMIT); \ }); \ } else { \ VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \ @@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) { scalar_t, typename vllm::PackedTypeConverter::Type, \ KERNEL, \ PACKED_KERNEL::Type>, \ - ACT_FIRST, true, false><<>>( \ - out.data_ptr(), input.data_ptr(), d); \ + ACT_FIRST, true, HAS_CLAMP, false><<>>( \ + out.data_ptr(), input.data_ptr(), d, LIMIT); \ }); \ } \ } else { \ @@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) { scalar_t, typename vllm::PackedTypeConverter::Type, \ KERNEL, \ PACKED_KERNEL::Type>, \ - ACT_FIRST, false><<>>( \ - out.data_ptr(), input.data_ptr(), d); \ + ACT_FIRST, false, HAS_CLAMP><<>>( \ + out.data_ptr(), input.data_ptr(), d, LIMIT); \ }); \ } @@ -206,7 +256,14 @@ 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); + 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); } void mul_and_silu(torch::Tensor& out, // [..., d] @@ -215,21 +272,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, false, 0.0f); } 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); + true, false, 0.0f); } 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); + LAUNCH_ACTIVATION_GATE_KERNEL( + vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f); } namespace vllm { diff --git a/csrc/async_util.cuh b/csrc/async_util.cuh new file mode 100644 index 000000000000..392d78c53fdb --- /dev/null +++ b/csrc/async_util.cuh @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +namespace vllm { +namespace cuda_async { + +__device__ __forceinline__ void cp_async_shared_global_16_cg( + void* smem_ptr, const void* glob_ptr) { +#if defined(USE_ROCM) + *reinterpret_cast(smem_ptr) = *reinterpret_cast(glob_ptr); +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem), "l"(glob_ptr)); +#elif defined(__CUDA_ARCH__) + *reinterpret_cast(smem_ptr) = *reinterpret_cast(glob_ptr); +#else + (void)smem_ptr; + (void)glob_ptr; +#endif +} + +__device__ __forceinline__ void cp_async_shared_global_ca(void* smem_ptr, + const void* glob_ptr, + int size_bytes) { +#if defined(USE_ROCM) + if (size_bytes == 4) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else if (size_bytes == 8) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if (size_bytes == 4) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" + : + : "r"(smem), "l"(glob_ptr)); + } else if (size_bytes == 8) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" + : + : "r"(smem), "l"(glob_ptr)); + } else { + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" + : + : "r"(smem), "l"(glob_ptr)); + } +#elif defined(__CUDA_ARCH__) + if (size_bytes == 4) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else if (size_bytes == 8) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } +#else + (void)smem_ptr; + (void)glob_ptr; + (void)size_bytes; +#endif +} + +__device__ __forceinline__ void cp_async_commit_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM) + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +template +__device__ __forceinline__ void cp_async_wait_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM) + asm volatile("cp.async.wait_group %0;\n" : : "n"(n)); +#endif +} + +} // namespace cuda_async +} // namespace vllm diff --git a/csrc/attention/dtype_fp8.cuh b/csrc/attention/dtype_fp8.cuh index e714e321b0be..1afec3c39976 100644 --- a/csrc/attention/dtype_fp8.cuh +++ b/csrc/attention/dtype_fp8.cuh @@ -17,6 +17,22 @@ enum class Fp8KVCacheDataType { kFp8E5M2 = 2, }; +inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( + const std::string& dtype_str) { + // dtype_str refers to CacheDType at vllm.config.cache.CacheDType + if (dtype_str == "auto" || dtype_str == "float16" || + dtype_str == "bfloat16") { + // unquantized kv cache + return Fp8KVCacheDataType::kAuto; + } else if (dtype_str == "fp8" || dtype_str == "fp8_ds_mla" || + dtype_str == "fp8_e4m3") { + return Fp8KVCacheDataType::kFp8E4M3; + } else if (dtype_str == "fp8_e5m2") { + return Fp8KVCacheDataType::kFp8E5M2; + } + TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); +} + // fp8 vector types for quantization of kv cache template <> struct Vec { diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index a7e5fdabf648..7e456d32598b 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -104,37 +104,49 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, static_assert(sizeof(CUdeviceptr) == sizeof(int64_t)); static_assert(sizeof(size_t) == sizeof(int64_t)); #if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 - CUmemcpyAttributes attr = {}; - attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; - size_t attrs_idx = 0; - #if defined(CUDA_VERSION) && CUDA_VERSION >= 13000 - CUresult result = cuMemcpyBatchAsync( - reinterpret_cast(dst_data), - reinterpret_cast(src_data), - reinterpret_cast(size_data), static_cast(n), &attr, - &attrs_idx, 1, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed with error ", - result); - #else - size_t fail_idx = 0; - CUresult result = cuMemcpyBatchAsync( - reinterpret_cast(dst_data), - reinterpret_cast(src_data), - reinterpret_cast(size_data), static_cast(n), &attr, - &attrs_idx, 1, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); - #endif -#else - // Fallback for CUDA < 12.8 and ROCm: individual async copies. - // cudaMemcpyDefault lets the driver infer direction from pointer types. - for (int64_t i = 0; i < n; i++) { - cudaMemcpyAsync(reinterpret_cast(dst_data[i]), - reinterpret_cast(src_data[i]), - static_cast(size_data[i]), cudaMemcpyDefault, - stream); - } + // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that + // binaries compiled with CUDA 12.8+ still work on older drivers, and + // we avoid the CUDA 13.0 header remapping (#define to _v2 signature). + // The function pointer is cached after the first call. + using BatchFn = + CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t, + CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream); + static BatchFn batch_fn = []() -> BatchFn { + CUdriverProcAddressQueryResult sym_status; + void* fn_ptr = nullptr; + CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080, + CU_GET_PROC_ADDRESS_DEFAULT, &sym_status); + if (res != CUDA_SUCCESS || fn_ptr == nullptr) { + return nullptr; + } + return reinterpret_cast(fn_ptr); + }(); + + if (batch_fn != nullptr) { + CUmemcpyAttributes attr = {}; + attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; + size_t attrs_idx = 0; + size_t fail_idx = 0; + CUresult result = batch_fn(reinterpret_cast(dst_data), + reinterpret_cast(src_data), + reinterpret_cast(size_data), + static_cast(n), &attr, &attrs_idx, 1, + &fail_idx, static_cast(stream)); + TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", + fail_idx, " with error ", result); + } else #endif + { + // Fallback for CUDA < 12.8, older drivers, and ROCm: + // individual async copies. + // cudaMemcpyDefault lets the driver infer direction from pointer types. + for (int64_t i = 0; i < n; i++) { + cudaMemcpyAsync(reinterpret_cast(dst_data[i]), + reinterpret_cast(src_data[i]), + static_cast(size_data[i]), cudaMemcpyDefault, + stream); + } + } } namespace vllm { @@ -587,6 +599,11 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE; // Find batch index within a block __shared__ int batch_idx[BLOCK_Y_SIZE]; + if (threadIdx.x == 0) { + batch_idx[threadIdx.y] = -1; + } + __syncthreads(); + for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); iter++) { int tid = iter * blockDim.x + threadIdx.x; @@ -599,16 +616,18 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( } } -#ifndef USE_ROCM - __syncwarp(); -#endif + __syncthreads(); - if (head_idx >= head_dim || token_idx >= num_tokens) { + // num_tokens may be an allocation upper bound when Python avoids a D2H sync. + // Only tokens covered by the exact device-side cu_seq_lens are valid to + // gather. + const int batch = batch_idx[threadIdx.y]; + if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) { return; } - const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]]; - const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks + - inbatch_seq_idx / cache_block_size]; + const int inbatch_seq_idx = token_idx - cu_seq_lens[batch]; + const int block_idx = + block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size]; const int64_t src_block_offset = block_idx * block_stride; const int64_t cache_inblock_offset = (inbatch_seq_idx % cache_block_size) * head_dim + head_idx; @@ -712,6 +731,28 @@ void reshape_and_cache_flash( int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + if (kv_cache_dtype == "nvfp4") { +#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) + // NVFP4 dispatch is compiled separately for SM100+. + extern void reshape_and_cache_nvfp4_dispatch( + torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache, + torch::Tensor & value_cache, torch::Tensor & slot_mapping, + torch::Tensor & k_scale, torch::Tensor & v_scale); + reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, + slot_mapping, k_scale, v_scale); + return; +#else + TORCH_CHECK(false, + "NVFP4 KV cache requires SM100+ (Blackwell). " + "Please rebuild vllm with a Blackwell-compatible CUDA target."); +#endif + } + + // Original FP8/auto path. int block_size = key_cache.size(1); int64_t key_stride = key.stride(0); @@ -729,8 +770,6 @@ void reshape_and_cache_flash( dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_size, 512)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE_FLASH); @@ -1458,6 +1497,9 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2"); TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); + TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half || + ql_nope.scalar_type() == at::ScalarType::BFloat16, + "ql_nope must be float16 or bfloat16 dtype"); if (num_tokens == 0) return; @@ -1469,7 +1511,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { + VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { vllm::ConcatMLAQKernel<<>>( q_out.data_ptr(), ql_nope.data_ptr(), q_pe.data_ptr(), num_tokens, num_heads, q_out.stride(0), diff --git a/csrc/cpu/activation_lut_bf16.cpp b/csrc/cpu/activation_lut_bf16.cpp new file mode 100644 index 000000000000..0ff2567e1ee8 --- /dev/null +++ b/csrc/cpu/activation_lut_bf16.cpp @@ -0,0 +1,71 @@ +#include "cpu_types.hpp" + +#include +#include +#include +#include + +#include +#include +#include + +constexpr uint32_t ActivationLutSize = 1u << 16; + +at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); } + +void maybe_init_activation_lut_bf16( + uint16_t* lut, std::once_flag& once, + at::Tensor (*activation)(const at::Tensor&)) { + std::call_once(once, [&]() { + auto lut_input = + at::empty({static_cast(ActivationLutSize)}, + at::TensorOptions().device(at::kCPU).dtype(at::kFloat)); + auto* lut_input_ptr = lut_input.data_ptr(); +#pragma omp parallel for + for (uint32_t i = 0; i < ActivationLutSize; ++i) { + lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast(i)); + } + + auto lut_output = activation(lut_input); + const auto* lut_output_ptr = lut_output.data_ptr(); +#pragma omp parallel for + for (uint32_t i = 0; i < ActivationLutSize; ++i) { + lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]); + } + }); +} + +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const uint16_t* lut, const char* op_name) { + TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name, + ": input must be bfloat16"); + TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name, + ": out must be bfloat16"); + TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous"); + TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous"); + + const auto* src = + reinterpret_cast(input.data_ptr()); + auto* dst = reinterpret_cast(out.data_ptr()); + const int64_t n = input.numel(); + + CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl) +#pragma omp parallel for + for (int64_t i = 0; i < n; ++i) { + dst[i] = lut[src[i]]; + } + CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl) +} + +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const std::string& activation) { + if (activation == "gelu") { + static std::array lut{}; + static std::once_flag once; + maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference); + activation_lut_bf16(out, input, lut.data(), "gelu_lut"); + return; + } + + TORCH_CHECK(false, "Unsupported activation: ", activation); +} diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index c73b62ecdec9..9be45a3efcef 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -61,8 +61,23 @@ #endif #ifdef __aarch64__ - // Implementation copied from Arm Optimized Routines (expf AdvSIMD) + // Implementation of neon_expf copied from Arm Optimized Routines (expf + // AdvSIMD) // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c + // + // Additional fast exponential intended for cases where outputs will be + // downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP + // for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] & + // clamps inputs outside this range to 0 / inf. Implementation is similar to + // exp_u20, but: + // - uses a third degree polynomial approximation for exp(r) instead of a + // fifth degree one, with coefficients re-tuned. + // - does not split natural log (ln) into high / low parts + // - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f + // exp(x) = 2^n (exp(r)) + // r = x - n*ln2, with n = round(x/ln2) + // exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r) + // n = round(x / ln2), r = x - n*ln2 #include #define DEFINE_FAST_EXP \ const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \ @@ -106,7 +121,38 @@ result.val[2] = neon_expf(vec.reg.val[2]); \ result.val[3] = neon_expf(vec.reg.val[3]); \ return vec_op::FP32Vec16(result); \ - }; + }; \ + const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \ + const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \ + constexpr float ln2 = 0x1.62e43p-1f; \ + constexpr float f_c2 = 0x1.5592ecp-3f; \ + const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \ + auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \ + always_inline)) { \ + const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \ + const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \ + float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \ + float32x4_t r = vfmsq_n_f32(values, n, ln2); \ + uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \ + float32x4_t r2 = vmulq_f32(r, r); \ + float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \ + float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \ + float32x4_t p = vfmaq_f32(s, q, r2); \ + float32x4_t y = \ + vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \ + y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \ + y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \ + return y; \ + }; \ + auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { \ + float32x4x4_t result; \ + result.val[0] = neon_expf_f16(vec.reg.val[0]); \ + result.val[1] = neon_expf_f16(vec.reg.val[1]); \ + result.val[2] = neon_expf_f16(vec.reg.val[2]); \ + result.val[3] = neon_expf_f16(vec.reg.val[3]); \ + return vec_op::FP32Vec16(result); \ + }; #endif // __aarch64__ diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index c15799fa950d..c1974bfd0a51 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -147,6 +147,9 @@ struct AttentionMetadata { case ISA::NEON: ss << "NEON, "; break; + case ISA::VXE: + ss << "VXE, "; + break; } ss << "workitem_group_num: " << workitem_group_num << ", reduction_item_num: " << reduction_item_num @@ -1149,7 +1152,11 @@ class AttentionMainLoop { bool use_sink) { #ifdef DEFINE_FAST_EXP DEFINE_FAST_EXP + bool constexpr IsReducedPrecision = + std::is_same_v || + std::is_same_v; #endif + using prob_buffer_vec_t = typename VecTypeTrait::vec_t; static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t)); @@ -1198,8 +1205,17 @@ class AttentionMainLoop { vec = vec - max_vec; // compute exp -#ifdef DEFINE_FAST_EXP - vec = fast_exp(vec); + +#if defined(DEFINE_FAST_EXP) + #ifdef __aarch64__ + if constexpr (IsReducedPrecision) { + vec = fast_exp_f16(vec); + } else + #endif + { + vec = fast_exp(vec); + } + prob_buffer_vec_t output_vec(vec); output_vec.save(curr_prob_buffer_iter); #else @@ -1255,7 +1271,11 @@ class AttentionMainLoop { int32_t kv_tile_token_num, float softcap_scale) { #ifdef DEFINE_FAST_EXP DEFINE_FAST_EXP + bool constexpr IsReducedPrecision = + std::is_same_v || + std::is_same_v; #endif + float inv_softcap_scale = 1.0 / softcap_scale; vec_op::FP32Vec16 softcap_scale_vec(softcap_scale); vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale); @@ -1269,8 +1289,15 @@ class AttentionMainLoop { vec_op::FP32Vec16 vec(curr_logits_buffer_iter); vec = vec * inv_softcap_scale_vec; -#ifdef DEFINE_FAST_EXP - vec = fast_exp(vec); +#if defined(DEFINE_FAST_EXP) + #ifdef __aarch64__ + if constexpr (IsReducedPrecision) { + vec = fast_exp_f16(vec); + } else + #endif + { + vec = fast_exp(vec); + } vec_op::FP32Vec16 inv_vec = ones_vec / vec; vec = (vec - inv_vec) / (vec + inv_vec); #else diff --git a/csrc/cpu/cpu_types_riscv.hpp b/csrc/cpu/cpu_types_riscv.hpp index 910ee5c11331..e617d98dd002 100644 --- a/csrc/cpu/cpu_types_riscv.hpp +++ b/csrc/cpu/cpu_types_riscv.hpp @@ -1,832 +1,25 @@ #ifndef CPU_TYPES_RISCV_HPP #define CPU_TYPES_RISCV_HPP -#include -#include -#include -#include -#include -#include -#include +// RISC-V Vector (RVV) CPU type definitions for vLLM. +// +// Supports multiple VLENs via compile-time dispatch. The compiler defines +// __riscv_v_min_vlen from the zvlb extension in -march. The defs header +// maps VLEN to the correct LMUL suffixes, and the impl header provides +// VLEN-independent class implementations. +// +// To add support for a new VLEN, add the LMUL mapping in +// cpu_types_riscv_defs.hpp (the impl header needs no changes). -// ============================================================================ -// Vector Register Type Definitions (VLEN=128 bits) -// ============================================================================ - -typedef vfloat16m1_t fixed_vfloat16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vfloat16m2_t fixed_vfloat16m2_t - __attribute__((riscv_rvv_vector_bits(256))); - -typedef vfloat32m1_t fixed_vfloat32m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vfloat32m2_t fixed_vfloat32m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vfloat32m4_t fixed_vfloat32m4_t - __attribute__((riscv_rvv_vector_bits(512))); -typedef vfloat32m8_t fixed_vfloat32m8_t - __attribute__((riscv_rvv_vector_bits(1024))); - -typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256))); -typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512))); - -typedef vuint16m1_t fixed_vuint16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vuint16m2_t fixed_vuint16m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vuint16m4_t fixed_vuint16m4_t - __attribute__((riscv_rvv_vector_bits(512))); - -#ifdef RISCV_BF16_SUPPORT -typedef vbfloat16m1_t fixed_vbfloat16m1_t - __attribute__((riscv_rvv_vector_bits(128))); -typedef vbfloat16m2_t fixed_vbfloat16m2_t - __attribute__((riscv_rvv_vector_bits(256))); -typedef vbfloat16m4_t fixed_vbfloat16m4_t - __attribute__((riscv_rvv_vector_bits(512))); -#endif - -namespace vec_op { - -#ifdef RISCV_BF16_SUPPORT - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) -#else - #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) +#ifndef __riscv_vector + #error "cpu_types_riscv.hpp included in a non-RVV translation unit" #endif -#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) - -#define FORCE_INLINE __attribute__((always_inline)) inline - -namespace { -template -constexpr void unroll_loop_item(std::integer_sequence, F&& f) { - (f(std::integral_constant{}), ...); -}; -} // namespace - -template >> -constexpr void unroll_loop(F&& f) { - unroll_loop_item(std::make_integer_sequence{}, std::forward(f)); -} - -template -struct Vec { - constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }; -}; - -struct FP32Vec8; -struct FP32Vec16; - -// ============================================================================ -// FP16 Implementation -// ============================================================================ - -struct FP16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat16m1_t reg; - - explicit FP16Vec8(const void* ptr) - : reg(__riscv_vle16_v_f16m1(static_cast(ptr), - VEC_ELEM_NUM)) {}; - - explicit FP16Vec8(const FP32Vec8&); - - void save(void* ptr) const { - __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(_Float16); - __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg, - VEC_ELEM_NUM); - } -}; - -struct FP16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat16m2_t reg; - - explicit FP16Vec16(const void* ptr) - : reg(__riscv_vle16_v_f16m2(static_cast(ptr), - VEC_ELEM_NUM)) {}; - - explicit FP16Vec16(const FP32Vec16& vec); - - void save(void* ptr) const { - __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(_Float16); - __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg, - VEC_ELEM_NUM); - } -}; - -// ============================================================================ -// BF16 Implementation -// ============================================================================ - -#ifdef RISCV_BF16_SUPPORT - -FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) { - return __riscv_vreinterpret_v_bf16m1_u16m1(v); -} -FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) { - return __riscv_vreinterpret_v_bf16m2_u16m2(v); -} -FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) { - return __riscv_vreinterpret_v_bf16m4_u16m4(v); -} - -struct BF16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vbfloat16m1_t reg; - - explicit BF16Vec8(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {}; - explicit BF16Vec8(const FP32Vec8&); - - void save(void* ptr) const { - __riscv_vse16_v_u16m1(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m1(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m1(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -struct BF16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vbfloat16m2_t reg; - - explicit BF16Vec16(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {}; - explicit BF16Vec16(const FP32Vec16&); - - void save(void* ptr) const { - __riscv_vse16_v_u16m2(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m2(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m2(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -struct BF16Vec32 : public Vec { - constexpr static int VEC_ELEM_NUM = 32; - fixed_vbfloat16m4_t reg; - - explicit BF16Vec32(const void* ptr) - : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4( - reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; - - explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {}; - - explicit BF16Vec32(const BF16Vec8& v) { - fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg); - fixed_vuint16m4_t u16_combined = - __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val); - reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined); - }; - - void save(void* ptr) const { - __riscv_vse16_v_u16m4(reinterpret_cast(ptr), bf16_to_u16(reg), - VEC_ELEM_NUM); - } - void save(void* ptr, int elem_num) const { - __riscv_vse16_v_u16m4(reinterpret_cast(ptr), bf16_to_u16(reg), - elem_num); - } - void save_strided(void* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - __riscv_vsse16_v_u16m4(reinterpret_cast(ptr), byte_stride, - bf16_to_u16(reg), VEC_ELEM_NUM); - } -}; - -#else -// ============================================================================ -// BF16 Fallback Implementation (FP32 Simulation) -// ============================================================================ - -struct BF16Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat32m2_t reg_fp32; - explicit BF16Vec8(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[8]; - for (int i = 0; i < 8; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8); - } - explicit BF16Vec8(const FP32Vec8&); - void save(void* ptr) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 8; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save(void* ptr, int elem_num) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[8]; - __riscv_vse32_v_f32m2(tmp, reg_fp32, 8); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 8; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; - -struct BF16Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat32m4_t reg_fp32; - explicit BF16Vec16(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[16]; - for (int i = 0; i < 16; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16); - } - explicit BF16Vec16(const FP32Vec16&); - void save(void* ptr) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 16; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save(void* ptr, int elem_num) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[16]; - __riscv_vse32_v_f32m4(tmp, reg_fp32, 16); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 16; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; - -struct BF16Vec32 : public Vec { - constexpr static int VEC_ELEM_NUM = 32; - fixed_vfloat32m8_t reg_fp32; - - explicit BF16Vec32(const void* ptr) { - const uint16_t* u16 = static_cast(ptr); - float tmp[32]; - for (int i = 0; i < 32; ++i) { - uint32_t v = static_cast(u16[i]) << 16; - std::memcpy(&tmp[i], &v, 4); - } - reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32); - } - - explicit BF16Vec32(const BF16Vec8& v) { - float tmp_small[8]; - __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8); - float tmp_large[32]; - for (int i = 0; i < 4; ++i) { - std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float)); - } - reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32); - } - - void save(void* ptr) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < 32; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - - void save(void* ptr, int elem_num) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint16_t* u16 = static_cast(ptr); - for (int i = 0; i < elem_num; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - u16[i] = static_cast(v >> 16); - } - } - - void save_strided(void* ptr, ptrdiff_t stride) const { - float tmp[32]; - __riscv_vse32_v_f32m8(tmp, reg_fp32, 32); - uint8_t* u8 = static_cast(ptr); - ptrdiff_t byte_stride = stride * sizeof(uint16_t); - for (int i = 0; i < 32; ++i) { - uint32_t v; - std::memcpy(&v, &tmp[i], 4); - uint16_t val = static_cast(v >> 16); - *reinterpret_cast(u8 + i * byte_stride) = val; - } - } -}; -#endif - -// ============================================================================ -// FP32 Implementation -// ============================================================================ - -struct FP32Vec4 : public Vec { - constexpr static int VEC_ELEM_NUM = 4; - fixed_vfloat32m1_t reg; - explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec4(const float* ptr) - : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {}; - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}; - void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m1(ptr, reg, elem_num); - } -}; - -struct FP32Vec8 : public Vec { - constexpr static int VEC_ELEM_NUM = 8; - fixed_vfloat32m2_t reg; - - explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(const float* ptr) - : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {}; - explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; - explicit FP32Vec8(const FP16Vec8& v) - : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(fixed_vfloat16m1_t v) - : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {}; - -#ifdef RISCV_BF16_SUPPORT - explicit FP32Vec8(fixed_vbfloat16m1_t v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec8(const BF16Vec8& v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {}; -#else - explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {}; +#ifndef __riscv_v_min_vlen + #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvlb" #endif - float reduce_sum() const { - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - FP32Vec8 operator*(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator+(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator-(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 operator/(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - - FP32Vec8 min(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 max(const FP32Vec8& b) const { - return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec8 abs() const { - return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM)); - } - - FP32Vec8 min(const FP32Vec8& b, int elem_num) const { - return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num)); - } - FP32Vec8 max(const FP32Vec8& b, int elem_num) const { - return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num)); - } - - FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const { - fixed_vfloat32m2_t temp = - __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM)); - } - - void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m2(ptr, reg, elem_num); - } - void save_strided(float* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(float); - __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM); - } - - FP32Vec8 exp() const { - const float inv_ln2 = 1.44269504088896341f; - fixed_vfloat32m2_t x_scaled = - __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM); - fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM); - fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM); - - fixed_vfloat32m2_t r = - __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM); - - fixed_vfloat32m2_t poly = - __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM); - - fixed_vint32m2_t biased_exp = - __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM); - biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM); - fixed_vint32m2_t exponent_bits = - __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM); - fixed_vfloat32m2_t scale = - __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits); - - return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM)); - } - - FP32Vec8 tanh() const { - fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2( - __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM); - fixed_vfloat32m2_t x2 = - __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM); - FP32Vec8 exp_val = FP32Vec8(x2).exp(); - fixed_vfloat32m2_t num = - __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM); - fixed_vfloat32m2_t den = - __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM)); - } - - FP32Vec8 er() const { - const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, - a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; - fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM); - - fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2( - __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM); - t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM); - - fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a4, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a3, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a2, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM), - a1, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM); - - fixed_vfloat32m2_t exp_val = - FP32Vec8(__riscv_vfneg_v_f32m2( - __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM), - VEC_ELEM_NUM)) - .exp() - .reg; - fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2( - __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f, - VEC_ELEM_NUM); - - vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM); - return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM)); - } -}; - -struct FP32Vec16 : public Vec { - constexpr static int VEC_ELEM_NUM = 16; - fixed_vfloat32m4_t reg; - - explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(const float* ptr) - : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {}; - explicit FP32Vec16(const FP32Vec8& data) - : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {}; - explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; - explicit FP32Vec16(const FP16Vec16& v); - -#ifdef RISCV_BF16_SUPPORT - explicit FP32Vec16(fixed_vbfloat16m2_t v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {}; - explicit FP32Vec16(const BF16Vec16& v) - : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {}; -#else - explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {}; -#endif - - FP32Vec16 operator+(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator-(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator*(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 operator/(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - - FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM)); - } - - float reduce_sum() const { - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - float reduce_max() const { - fixed_vfloat32m1_t scalar = - __riscv_vfmv_s_f_f32m1(std::numeric_limits::lowest(), 1); - scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - float reduce_min() const { - fixed_vfloat32m1_t scalar = - __riscv_vfmv_s_f_f32m1(std::numeric_limits::max(), 1); - scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - } - - template - float reduce_sub_sum(int idx) { - static_assert(VEC_ELEM_NUM % group_size == 0); - const int start = idx * group_size; - vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM); - vbool8_t mask = __riscv_vmand_mm_b8( - __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM), - __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM), - VEC_ELEM_NUM); - fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); - scalar = - __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM); - return __riscv_vfmv_f_s_f32m1_f32(scalar); - }; - - FP32Vec16 max(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 min(const FP32Vec16& b) const { - return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM)); - } - FP32Vec16 abs() const { - return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM)); - } - - FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { - return FP32Vec16(__riscv_vfmin_vv_f32m4( - max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM), - VEC_ELEM_NUM)); - } - - void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); } - void save(float* ptr, int elem_num) const { - __riscv_vse32_v_f32m4(ptr, reg, elem_num); - } - void save_strided(float* ptr, ptrdiff_t stride) const { - ptrdiff_t byte_stride = stride * sizeof(float); - __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM); - } - - FP32Vec16 exp() const { - const float inv_ln2 = 1.44269504088896341f; - fixed_vfloat32m4_t x_scaled = - __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM); - fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM); - fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM); - fixed_vfloat32m4_t r = - __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM); - - fixed_vfloat32m4_t poly = - __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.009618129107628f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.055504108664821f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.240226506959101f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 0.693147180559945f, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM); - - fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4( - __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM); - fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4( - __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM)); - - return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM)); - } - - FP32Vec16 tanh() const { - fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4( - __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM); - FP32Vec16 exp_val = - FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp(); - return FP32Vec16(__riscv_vfdiv_vv_f32m4( - __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), - __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM)); - } - - FP32Vec16 er() const { - const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, - a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; - fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM); - fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4( - __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM), - 1.0f, VEC_ELEM_NUM); - - fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a4, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a3, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a2, VEC_ELEM_NUM); - poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM), - a1, VEC_ELEM_NUM); - poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM); - - fixed_vfloat32m4_t exp_val = - FP32Vec16(__riscv_vfneg_v_f32m4( - __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM), - VEC_ELEM_NUM)) - .exp() - .reg; - fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4( - __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f, - VEC_ELEM_NUM); - - vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM); - return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM)); - } -}; - -// ============================================================================ -// Type Traits & Global Helpers -// ============================================================================ - -template -struct VecType { - using vec_type = void; - using vec_t = void; -}; - -template -using vec_t = typename VecType::vec_type; - -template <> -struct VecType { - using vec_type = FP32Vec8; - using vec_t = FP32Vec8; -}; -template <> -struct VecType { - using vec_type = FP16Vec8; - using vec_t = FP16Vec8; -}; -template <> -struct VecType { - using vec_type = BF16Vec8; - using vec_t = BF16Vec8; -}; - -template -void storeFP32(float v, T* ptr) { - *ptr = v; -} -template <> -inline void storeFP32(float v, c10::Half* ptr) { - *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v); -} - -inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { - reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM); -} -inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { - reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM); -} -inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { - reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM); -} -inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) { - acc = acc.fma(a, b); -} - -#ifdef RISCV_BF16_SUPPORT -template <> -inline void storeFP32(float v, c10::BFloat16* ptr) { - *ptr = static_cast<__bf16>(v); -}; -inline BF16Vec8::BF16Vec8(const FP32Vec8& v) - : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {}; -inline BF16Vec16::BF16Vec16(const FP32Vec16& v) - : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {}; -#else -template <> -inline void storeFP32(float v, c10::BFloat16* ptr) { - uint32_t val; - std::memcpy(&val, &v, 4); - *reinterpret_cast(ptr) = static_cast(val >> 16); -} -inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {} -inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {} -#endif - -inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); } - -} // namespace vec_op - -#ifndef CPU_KERNEL_GUARD_IN - #define CPU_KERNEL_GUARD_IN(NAME) -#endif - -#ifndef CPU_KERNEL_GUARD_OUT - #define CPU_KERNEL_GUARD_OUT(NAME) -#endif +#include "cpu_types_riscv_defs.hpp" +#include "cpu_types_riscv_impl.hpp" -#endif // CPU_TYPES_RISCV_HPP \ No newline at end of file +#endif // CPU_TYPES_RISCV_HPP diff --git a/csrc/cpu/cpu_types_riscv_defs.hpp b/csrc/cpu/cpu_types_riscv_defs.hpp new file mode 100644 index 000000000000..04ce2728f032 --- /dev/null +++ b/csrc/cpu/cpu_types_riscv_defs.hpp @@ -0,0 +1,98 @@ +#ifndef CPU_TYPES_RISCV_DEFS_HPP +#define CPU_TYPES_RISCV_DEFS_HPP + +// VLEN-to-LMUL mapping for RISC-V Vector extension. +// +// LMUL_ expands to the LMUL suffix giving N total bits of vector data: +// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8 +// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4 + +#include + +#if __riscv_v_min_vlen == 128 + #define LMUL_128 m1 + #define LMUL_256 m2 + #define LMUL_512 m4 + #define LMUL_1024 m8 + #define BOOL_256 b16 + #define BOOL_512 b8 +#elif __riscv_v_min_vlen == 256 + #define LMUL_128 mf2 + #define LMUL_256 m1 + #define LMUL_512 m2 + #define LMUL_1024 m4 + #define BOOL_256 b32 + #define BOOL_512 b16 +#else + #error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen" +#endif + +// Token-paste helpers. +#define _RVV_P2(a, b) a##b +#define _RVV_P3(a, b, c) a##b##c +#define _RVV_P4(a, b, c, d) a##b##c##d +#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix) +#define RVVI(base, lmul) _RVV_P2(base, lmul) +#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix) +#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d) +// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16 +#define _RVV_PB(base, lmul, btype) base##lmul##_##btype +#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype) + +// ---- Semantic fixed-vector typedefs (named by element count) ---- + +// float16 +typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t + __attribute__((riscv_rvv_vector_bits(256))); + +// float32 +typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t + __attribute__((riscv_rvv_vector_bits(512))); +typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t + __attribute__((riscv_rvv_vector_bits(1024))); + +// int32 +typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t + __attribute__((riscv_rvv_vector_bits(512))); + +// uint16 +typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t + __attribute__((riscv_rvv_vector_bits(512))); + +// bfloat16 +#ifdef RISCV_BF16_SUPPORT +typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t + __attribute__((riscv_rvv_vector_bits(512))); +#endif + +// ---- Reduction accumulator type (always m1 = one register of f32) ---- +// Used for scalar reductions; only element [0] is meaningful. +typedef vfloat32m1_t rvv_f32_accum_t + __attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen))); + +// ---- Mask types for f32 elements ---- +#if __riscv_v_min_vlen == 128 +typedef vbool16_t rvv_mask_f32x8_t; +typedef vbool8_t rvv_mask_f32x16_t; +#elif __riscv_v_min_vlen == 256 +typedef vbool32_t rvv_mask_f32x8_t; +typedef vbool16_t rvv_mask_f32x16_t; +#endif + +#endif // CPU_TYPES_RISCV_DEFS_HPP diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp new file mode 100644 index 000000000000..3952a811c2c5 --- /dev/null +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -0,0 +1,905 @@ +#ifndef CPU_TYPES_RISCV_IMPL_HPP +#define CPU_TYPES_RISCV_IMPL_HPP + +// Shared implementation of RVV vector-type wrapper classes. +// This file is VLEN-independent: it uses the semantic type names and +// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp. +// +// DO NOT include this file directly; include cpu_types_riscv.hpp instead. + +#include +#include +#include +#include +#include +#include +namespace vec_op { + +// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT +// is defined, otherwise via the FP32-simulation fallback path. +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define FORCE_INLINE __attribute__((always_inline)) inline + +namespace { +template +constexpr void unroll_loop_item(std::integer_sequence, F&& f) { + (f(std::integral_constant{}), ...); +}; +} // namespace + +template >> +constexpr void unroll_loop(F&& f) { + unroll_loop_item(std::make_integer_sequence{}, std::forward(f)); +} + +template +struct Vec { + constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }; +}; + +struct FP32Vec8; +struct FP32Vec16; + +// ============================================================================ +// FP16 Implementation +// ============================================================================ + +struct FP16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp16x8_t reg; + + explicit FP16Vec8(const void* ptr) + : reg(RVVI(__riscv_vle16_v_f16, LMUL_128)( + static_cast(ptr), VEC_ELEM_NUM)) {}; + + explicit FP16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg, + VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg, + elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(_Float16); + RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), + byte_stride, reg, VEC_ELEM_NUM); + } +}; + +struct FP16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp16x16_t reg; + + explicit FP16Vec16(const void* ptr) + : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)( + static_cast(ptr), VEC_ELEM_NUM)) {}; + + explicit FP16Vec16(const FP32Vec16& vec); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg, + VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg, + elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(_Float16); + RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), + byte_stride, reg, VEC_ELEM_NUM); + } +}; + +// ============================================================================ +// BF16 Implementation +// ============================================================================ + +#ifdef RISCV_BF16_SUPPORT + +FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v); +} +FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v); +} +FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) { + return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v); +} + +struct BF16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_bf16x8_t reg; + + explicit BF16Vec8(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, + LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {}; + explicit BF16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +struct BF16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_bf16x16_t reg; + + explicit BF16Vec16(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, + LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {}; + explicit BF16Vec16(const FP32Vec16&); + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +struct BF16Vec32 : public Vec { + constexpr static int VEC_ELEM_NUM = 32; + fixed_bf16x32_t reg; + + explicit BF16Vec32(const void* ptr) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, + LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)( + reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; + + explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {}; + + explicit BF16Vec32(const BF16Vec8& v) { + fixed_u16x8_t u16_val = bf16_to_u16(v.reg); + fixed_u16x32_t u16_combined = + RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)( + u16_val, u16_val, u16_val, u16_val); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, + LMUL_512)(u16_combined); + }; + + void save(void* ptr) const { + RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + bf16_to_u16(reg), VEC_ELEM_NUM); + } + void save(void* ptr, int elem_num) const { + RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + bf16_to_u16(reg), elem_num); + } + void save_strided(void* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast(ptr), + byte_stride, bf16_to_u16(reg), + VEC_ELEM_NUM); + } +}; + +#else +// ============================================================================ +// BF16 Fallback Implementation (FP32 Simulation) +// ============================================================================ + +struct BF16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp32x8_t reg_fp32; + explicit BF16Vec8(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[8]; + for (int i = 0; i < 8; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8); + } + explicit BF16Vec8(const FP32Vec8&); + void save(void* ptr) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 8; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save(void* ptr, int elem_num) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 8; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; + +struct BF16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp32x16_t reg_fp32; + explicit BF16Vec16(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[16]; + for (int i = 0; i < 16; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16); + } + explicit BF16Vec16(const FP32Vec16&); + void save(void* ptr) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 16; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save(void* ptr, int elem_num) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[16]; + RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 16; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; + +struct BF16Vec32 : public Vec { + constexpr static int VEC_ELEM_NUM = 32; + fixed_fp32x32_t reg_fp32; + + explicit BF16Vec32(const void* ptr) { + const uint16_t* u16 = static_cast(ptr); + float tmp[32]; + for (int i = 0; i < 32; ++i) { + uint32_t v = static_cast(u16[i]) << 16; + std::memcpy(&tmp[i], &v, 4); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32); + } + + explicit BF16Vec32(const BF16Vec8& v) { + float tmp_small[8]; + RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8); + float tmp_large[32]; + for (int i = 0; i < 4; ++i) { + std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float)); + } + reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32); + } + + void save(void* ptr) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < 32; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + + void save(void* ptr, int elem_num) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint16_t* u16 = static_cast(ptr); + for (int i = 0; i < elem_num; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + u16[i] = static_cast(v >> 16); + } + } + + void save_strided(void* ptr, ptrdiff_t stride) const { + float tmp[32]; + RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32); + uint8_t* u8 = static_cast(ptr); + ptrdiff_t byte_stride = stride * sizeof(uint16_t); + for (int i = 0; i < 32; ++i) { + uint32_t v; + std::memcpy(&v, &tmp[i], 4); + uint16_t val = static_cast(v >> 16); + *reinterpret_cast(u8 + i * byte_stride) = val; + } + } +}; +#endif + +// ============================================================================ +// FP32 Implementation +// ============================================================================ + +struct FP32Vec4 : public Vec { + constexpr static int VEC_ELEM_NUM = 4; + fixed_fp32x4_t reg; + explicit FP32Vec4(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec4() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec4(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {}; + explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}; + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num); + } +}; + +struct FP32Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + fixed_fp32x8_t reg; + + explicit FP32Vec8(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec8() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {}; + explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; + explicit FP32Vec8(const FP16Vec8& v) + : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(fixed_fp16x8_t v) + : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + +#ifdef RISCV_BF16_SUPPORT + explicit FP32Vec8(fixed_bf16x8_t v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec8(const BF16Vec8& v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) { + }; +#else + explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {}; +#endif + + float reduce_sum() const { + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + FP32Vec8 operator*(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator+(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator-(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 operator/(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + + FP32Vec8 min(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 max(const FP32Vec8& b) const { + return FP32Vec8( + RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec8 abs() const { + return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM)); + } + + FP32Vec8 min(const FP32Vec8& b, int elem_num) const { + return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num)); + } + FP32Vec8 max(const FP32Vec8& b, int elem_num) const { + return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num)); + } + + FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const { + fixed_fp32x8_t temp = + RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM); + return FP32Vec8( + RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM)); + } + + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num); + } + void save_strided(float* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(float); + RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM); + } + + FP32Vec8 exp() const { + // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN. + // Without clamping, -inf * 0.0 = NaN in the final poly * scale step. + // Matches the clamping strategy used by x86 AVX-512 and ARM NEON. + constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN) + constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX) + fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)( + RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi, + VEC_ELEM_NUM); + + const float inv_ln2 = 1.44269504088896341f; + fixed_fp32x8_t x_scaled = + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM); + fixed_i32x8_t n_int = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM); + fixed_fp32x8_t n_float = + RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM); + + fixed_fp32x8_t r = + RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM); + + fixed_fp32x8_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM); + + fixed_i32x8_t biased_exp = + RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM); + biased_exp = + RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM); + fixed_i32x8_t exponent_bits = + RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM); + fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32, + LMUL_256)(exponent_bits); + + return FP32Vec8( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM)); + } + + FP32Vec8 tanh() const { + fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)( + RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f, + VEC_ELEM_NUM); + fixed_fp32x8_t x2 = + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM); + FP32Vec8 exp_val = FP32Vec8(x2).exp(); + fixed_fp32x8_t num = + RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM); + fixed_fp32x8_t den = + RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM); + return FP32Vec8( + RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM)); + } + + FP32Vec8 er() const { + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; + fixed_fp32x8_t abs_x = + RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM); + + fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM); + + fixed_fp32x8_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM); + + fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)( + abs_x, abs_x, VEC_ELEM_NUM), + VEC_ELEM_NUM)) + .exp() + .reg; + fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)( + RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)( + reg, 0.0f, VEC_ELEM_NUM); + return FP32Vec8( + RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM)); + } +}; + +struct FP32Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_fp32x16_t reg; + + explicit FP32Vec16(float v) + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec16() + : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(const float* ptr) + : reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {}; + explicit FP32Vec16(const FP32Vec8& data) + : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)( + data.reg, data.reg)) {}; + explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; + explicit FP32Vec16(const FP16Vec16& v); + +#ifdef RISCV_BF16_SUPPORT + explicit FP32Vec16(fixed_bf16x16_t v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {}; + explicit FP32Vec16(const BF16Vec16& v) + : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) { + }; +#else + explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {}; +#endif + + FP32Vec16 operator+(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator-(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator*(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 operator/(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + + FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM)); + } + + float reduce_sum() const { + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + float reduce_max() const { + rvv_f32_accum_t scalar = + __riscv_vfmv_s_f_f32m1(std::numeric_limits::lowest(), 1); + scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + float reduce_min() const { + rvv_f32_accum_t scalar = + __riscv_vfmv_s_f_f32m1(std::numeric_limits::max(), 1); + scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar, + VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + } + + template + float reduce_sub_sum(int idx) { + static_assert(VEC_ELEM_NUM % group_size == 0); + const int start = idx * group_size; + auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM); + rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)( + RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start, + VEC_ELEM_NUM), + RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)( + indices, start + group_size, VEC_ELEM_NUM), + VEC_ELEM_NUM); + rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1); + scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)( + mask, reg, scalar, VEC_ELEM_NUM); + return __riscv_vfmv_f_s_f32m1_f32(scalar); + }; + + FP32Vec16 max(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 min(const FP32Vec16& b) const { + return FP32Vec16( + RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); + } + FP32Vec16 abs() const { + return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM)); + } + + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { + return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)( + max_v.reg, + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM), + VEC_ELEM_NUM)); + } + + void save(float* ptr) const { + RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM); + } + void save(float* ptr, int elem_num) const { + RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num); + } + void save_strided(float* ptr, ptrdiff_t stride) const { + ptrdiff_t byte_stride = stride * sizeof(float); + RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM); + } + + FP32Vec16 exp() const { + // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN. + // Without clamping, -inf * 0.0 = NaN in the final poly * scale step. + // Matches the clamping strategy used by x86 AVX-512 and ARM NEON. + constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN) + constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX) + fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)( + RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi, + VEC_ELEM_NUM); + + const float inv_ln2 = 1.44269504088896341f; + fixed_fp32x16_t x_scaled = + RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM); + fixed_i32x16_t n_int = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM); + fixed_fp32x16_t n_float = + RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM); + fixed_fp32x16_t r = + RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM); + + fixed_fp32x16_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.009618129107628f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.055504108664821f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.240226506959101f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), + 0.693147180559945f, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)( + RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0, + VEC_ELEM_NUM); + fixed_fp32x16_t scale = + RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)( + RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM)); + + return FP32Vec16( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM)); + } + + FP32Vec16 tanh() const { + fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)( + RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f, + VEC_ELEM_NUM); + FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)( + x_clamped, 2.0f, VEC_ELEM_NUM)) + .exp(); + return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)( + RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM), + RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM), + VEC_ELEM_NUM)); + } + + FP32Vec16 er() const { + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f, + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f; + fixed_fp32x16_t abs_x = + RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM); + fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)( + RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM), + 1.0f, VEC_ELEM_NUM); + + fixed_fp32x16_t poly = + RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1, + VEC_ELEM_NUM); + poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM); + + fixed_fp32x16_t exp_val = + FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x, + VEC_ELEM_NUM), + VEC_ELEM_NUM)) + .exp() + .reg; + fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)( + RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f, + VEC_ELEM_NUM); + + rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)( + reg, 0.0f, VEC_ELEM_NUM); + return FP32Vec16( + RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM)); + } +}; + +// ============================================================================ +// Type Traits & Global Helpers +// ============================================================================ + +template +struct VecType { + using vec_type = void; + using vec_t = void; +}; + +template +using vec_t = typename VecType::vec_type; + +template <> +struct VecType { + using vec_type = FP32Vec8; + using vec_t = FP32Vec8; +}; +template <> +struct VecType { + using vec_type = FP16Vec8; + using vec_t = FP16Vec8; +}; +template <> +struct VecType { + using vec_type = BF16Vec8; + using vec_t = BF16Vec8; +}; + +template +void storeFP32(float v, T* ptr) { + *ptr = v; +} +template <> +inline void storeFP32(float v, c10::Half* ptr) { + *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v); +} + +inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { + reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM); +} +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { + reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM); +} +inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { + reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM); +} +inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) { + acc = acc.fma(a, b); +} + +#ifdef RISCV_BF16_SUPPORT +template <> +inline void storeFP32(float v, c10::BFloat16* ptr) { + *ptr = static_cast<__bf16>(v); +}; +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) + : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) { + }; +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) + : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) { + }; +#else +template <> +inline void storeFP32(float v, c10::BFloat16* ptr) { + uint32_t val; + std::memcpy(&val, &v, 4); + *reinterpret_cast(ptr) = static_cast(val >> 16); +} +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {} +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {} +#endif + +inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); } + +} // namespace vec_op + +#ifndef CPU_KERNEL_GUARD_IN + #define CPU_KERNEL_GUARD_IN(NAME) +#endif + +#ifndef CPU_KERNEL_GUARD_OUT + #define CPU_KERNEL_GUARD_OUT(NAME) +#endif + +#endif // CPU_TYPES_RISCV_IMPL_HPP diff --git a/csrc/cpu/pos_encoding.cpp b/csrc/cpu/pos_encoding.cpp index 74bb014cf39e..9f41e4e222bd 100644 --- a/csrc/cpu/pos_encoding.cpp +++ b/csrc/cpu/pos_encoding.cpp @@ -178,7 +178,12 @@ void rotary_embedding_gptj_impl( void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, std::optional key, int64_t head_size, - torch::Tensor& cos_sin_cache, bool is_neox) { + torch::Tensor& cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse) { + TORCH_CHECK(rope_dim_offset == 0, + "rope_dim_offset != 0 is not supported on CPU"); + TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU"); + int num_tokens = positions.numel(); int rot_dim = cos_sin_cache.size(1); int num_heads = query.size(-1) / head_size; diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp new file mode 100644 index 000000000000..a76b8bc69376 --- /dev/null +++ b/csrc/cpu/spec_decode_utils.cpp @@ -0,0 +1,409 @@ +#include "cpu_types.hpp" + +#include + +namespace cpu_utils { + +void eagle_prepare_inputs_padded_kernel_impl( + const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& valid_sampled_tokens_count, + const torch::Tensor& query_start_loc_gpu, + torch::Tensor& token_indices_to_sample, + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) { + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* valid_count_ptr = + valid_sampled_tokens_count.data_ptr(); + const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr(); + int32_t* indices_out_ptr = token_indices_to_sample.data_ptr(); + int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx; + int64_t num_valid_tokens = valid_count_ptr[req_idx]; + + int64_t num_rejected = 0; + if (num_draft_tokens > 0) { + num_rejected = num_draft_tokens + 1 - num_valid_tokens; + } + + int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1; + int32_t index_to_sample = q_last_tok_idx - num_rejected; + + indices_out_ptr[req_idx] = index_to_sample; + rejected_out_ptr[req_idx] = num_rejected; + } +} + +void eagle_prepare_next_token_padded_kernel_impl( + const torch::Tensor& sampled_token_ids, + const torch::Tensor& discard_request_mask, + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) { + const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr(); + const bool* discard_mask_ptr = discard_request_mask.data_ptr(); + const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr(); + int64_t* next_ids_out_ptr = next_token_ids.data_ptr(); + int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr(); + + const int64_t stride = sampled_token_ids.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride; + int64_t valid_count = 0; + int64_t last_valid_token = -1; + + for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) { + int64_t token = row_ptr[pos]; + if (token != -1 && token < vocab_size) { + valid_count++; + last_valid_token = token; + } + } + + bool discard = discard_mask_ptr[req_idx]; + if (discard) { + next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx]; + valid_count_out_ptr[req_idx] = 0; + } else { + next_ids_out_ptr[req_idx] = + (valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx]; + valid_count_out_ptr[req_idx] = valid_count; + } + } +} + +void eagle_step_slot_mapping_metadata_kernel_impl( + const torch::Tensor& positions, const torch::Tensor& block_table, + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, + torch::Tensor& out_slot_mapping, const int64_t block_size, + const int64_t max_model_len, const int64_t PAD_ID) { + const int64_t batch_size = positions.size(0); + const int64_t input_batch_size = out_slot_mapping.size(0); + + const int64_t* pos_ptr = positions.data_ptr(); + const int32_t* bt_ptr = block_table.data_ptr(); + int32_t* seq_lens_ptr = seq_lens.data_ptr(); + int64_t* out_clamped_ptr = out_clamped_positions.data_ptr(); + int64_t* out_slot_ptr = out_slot_mapping.data_ptr(); + + const int64_t bt_stride = block_table.stride(0); + const int64_t n_blocks_per_req = block_table.size(1); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) { + if (req_idx >= batch_size) { + out_slot_ptr[req_idx] = PAD_ID; + continue; + } + + int64_t position = pos_ptr[req_idx]; + int64_t new_position = position + 1; + bool exceeds_max = new_position >= max_model_len; + int64_t clamped_position = exceeds_max ? 0 : new_position; + + out_clamped_ptr[req_idx] = clamped_position; + + int64_t block_number = clamped_position / block_size; + block_number = std::min(block_number, n_blocks_per_req - 1); + int32_t block_id = bt_ptr[req_idx * bt_stride + block_number]; + int64_t slot_id = block_id * block_size + (clamped_position % block_size); + out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id; + + int32_t seq_len = seq_lens_ptr[req_idx]; + int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1); + new_seq_len = std::min(new_seq_len, static_cast(max_model_len)); + seq_lens_ptr[req_idx] = new_seq_len; + } +} + +void copy_and_expand_eagle_inputs_kernel_impl( + const torch::Tensor& target_token_ids, + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, + torch::Tensor& out_input_ids, torch::Tensor& out_positions, + torch::Tensor& out_is_rejected_token_mask, + torch::Tensor& out_is_masked_token_mask, + torch::Tensor& out_new_token_indices, + torch::Tensor& out_hidden_state_mapping, + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, + const int64_t padding_token_id, const int64_t parallel_drafting_token_id, + const int64_t total_input_tokens, + const int64_t num_padding_slots_per_request, const bool shift_input_ids) { + const int64_t num_reqs = query_end_loc.size(0); + + const int64_t* target_ids_ptr = target_token_ids.data_ptr(); + const int64_t* target_pos_ptr = target_positions.data_ptr(); + const int64_t* next_ids_ptr = next_token_ids.data_ptr(); + const int32_t* query_start_ptr = query_start_loc.data_ptr(); + const int32_t* query_end_ptr = query_end_loc.data_ptr(); + + int64_t* out_ids_ptr = out_input_ids.data_ptr(); + int64_t* out_pos_ptr = out_positions.data_ptr(); + bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr(); + bool* out_mask_ptr = out_is_masked_token_mask.data_ptr(); + int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr(); + int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int32_t q_start = query_start_ptr[req_idx]; + int32_t next_q_start = query_start_ptr[req_idx + 1]; + int32_t q_end = query_end_ptr[req_idx]; + + int64_t num_valid_tokens = + shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1); + int64_t input_offset = shift_input_ids ? 1 : 0; + + int64_t out_start = q_start + req_idx * (num_padding_slots_per_request - + (shift_input_ids ? 1 : 0)); + int64_t num_rejected = next_q_start - q_end - 1; + int64_t total_output_tokens = + num_valid_tokens + num_padding_slots_per_request + num_rejected; + + int64_t start_pos = target_pos_ptr[q_start]; + int64_t bonus_token = next_ids_ptr[req_idx]; + + for (int64_t j = 0; j < total_output_tokens; ++j) { + int64_t out_idx = out_start + j; + bool is_valid = j < num_valid_tokens; + bool is_bonus = j == num_valid_tokens; + bool is_parallel = (j > num_valid_tokens) && + (j < num_valid_tokens + num_padding_slots_per_request); + bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request; + + int64_t in_idx = + std::min(static_cast(q_start + input_offset + j), + total_input_tokens - 1); + + int64_t token_id = padding_token_id; + if (is_valid) + token_id = target_ids_ptr[in_idx]; + else if (is_bonus) + token_id = bonus_token; + else if (is_parallel) + token_id = parallel_drafting_token_id; + + out_ids_ptr[out_idx] = token_id; + out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j); + out_rej_mask_ptr[out_idx] = is_rejected; + out_mask_ptr[out_idx] = is_parallel; + + if (is_bonus || is_parallel) { + int64_t new_token_local_idx = j - num_valid_tokens; + int64_t new_token_out_idx = + req_idx * num_padding_slots_per_request + new_token_local_idx; + out_new_idx_ptr[new_token_out_idx] = out_idx; + } + } + + if (shift_input_ids) { + int64_t n_input = next_q_start - q_start; + for (int64_t j = 0; j < n_input; ++j) { + out_hidden_map_ptr[q_start + j] = out_start + j; + } + } + } +} + +void rejection_greedy_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, + const torch::Tensor& bonus_token_ids, + const std::optional& is_greedy, const int64_t max_spec_len) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const int64_t* target_argmax_ptr = target_argmax.data_ptr(); + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr(); + const bool* greedy_ptr = + is_greedy.has_value() ? is_greedy.value().data_ptr() : nullptr; + + const int64_t out_stride = output_token_ids.stride(0); + const int64_t bonus_stride = bonus_token_ids.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + if (greedy_ptr && !greedy_ptr[req_idx]) continue; + + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + bool rejected = false; + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t target_id = target_argmax_ptr[start_idx + pos]; + out_ptr[req_idx * out_stride + pos] = target_id; + + if (draft_ids_ptr[start_idx + pos] != target_id) { + rejected = true; + break; + } + } + + if (!rejected) { + out_ptr[req_idx * out_stride + num_draft_tokens] = + bonus_ids_ptr[req_idx * bonus_stride]; + } + } +} + +void rejection_random_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, + const torch::Tensor& recovered_token_ids, + const torch::Tensor& uniform_probs, + const std::optional& is_greedy, const int64_t max_spec_len, + const int64_t vocab_size, const bool no_draft_probs) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const float* draft_probs_ptr = + no_draft_probs ? nullptr : draft_probs.value().data_ptr(); + const float* target_probs_ptr = target_probs.data_ptr(); + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr(); + const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr(); + const float* uniform_probs_ptr = uniform_probs.data_ptr(); + const bool* greedy_ptr = + is_greedy.has_value() ? is_greedy.value().data_ptr() : nullptr; + + const int64_t out_stride = output_token_ids.stride(0); + const int64_t bonus_stride = bonus_token_ids.stride(0); + const int64_t target_stride = target_probs.stride(0); + const int64_t draft_probs_stride = + no_draft_probs ? 0 : draft_probs.value().stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + if (greedy_ptr && greedy_ptr[req_idx]) continue; + + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + bool rejected = false; + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t token_idx = start_idx + pos; + int64_t draft_id = draft_ids_ptr[token_idx]; + + float p = target_probs_ptr[token_idx * target_stride + draft_id]; + float q = + no_draft_probs + ? 1.0f + : draft_probs_ptr[token_idx * draft_probs_stride + draft_id]; + float uniform_p = uniform_probs_ptr[token_idx]; + + float ratio = (q > 0.0f) ? (p / q) : 0.0f; + + if (ratio >= uniform_p) { + out_ptr[req_idx * out_stride + pos] = draft_id; + } else { + out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx]; + rejected = true; + break; + } + } + + if (!rejected) { + out_ptr[req_idx * out_stride + num_draft_tokens] = + bonus_ids_ptr[req_idx * bonus_stride]; + } + } +} + +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& cu_num_tokens, + const int64_t replace_from, const int64_t replace_to) { + const int64_t batch_size = cu_num_tokens.size(0); + const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr(); + + int64_t* out_ptr = output.data_ptr(); + const int64_t* in_ptr = input.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1]; + int64_t end_idx = cu_tokens_ptr[req_idx]; + int64_t val = in_ptr[req_idx]; + + if (val == replace_from) { + val = replace_to; + } + + for (int64_t i = start_idx; i < end_idx; ++i) { + out_ptr[i] = val; + } + } +} + +void sample_recovered_tokens_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& inv_q, + const int64_t vocab_size, const bool no_draft_probs) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const float* draft_probs_ptr = + no_draft_probs ? nullptr : draft_probs.value().data_ptr(); + const float* target_probs_ptr = target_probs.data_ptr(); + const float* inv_q_ptr = inv_q.data_ptr(); + + const int64_t target_stride = target_probs.stride(0); + const int64_t draft_probs_stride = + no_draft_probs ? 0 : draft_probs.value().stride(0); + const int64_t inv_q_stride = inv_q.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride; + + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t token_idx = start_idx + pos; + int64_t draft_id = draft_ids_ptr[token_idx]; + + const float* token_target_probs = + target_probs_ptr + token_idx * target_stride; + const float* token_draft_probs = + no_draft_probs ? nullptr + : (draft_probs_ptr + token_idx * draft_probs_stride); + + int64_t best_id = 0; + float best_val = -1.0f; + + for (int64_t v = 0; v < vocab_size; ++v) { + float prob = token_target_probs[v]; + if (no_draft_probs) { + if (v == draft_id) prob = 0.0f; + } else { + float diff = prob - token_draft_probs[v]; + prob = diff > 0.0f ? diff : 0.0f; + } + + float val = prob * req_inv_q[v]; + if (val > best_val) { + best_val = val; + best_id = v; + } + } + out_ptr[token_idx] = best_id; + } + } +} + +} // namespace cpu_utils diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index fbc9c65241b6..bd57c9188886 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional bias); +void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, + const std::string& activation); + torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, const int64_t num_heads_kv, const int64_t head_dim, @@ -138,6 +141,63 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, torch::Tensor slot_mapping, const int64_t block_size); +void init_cpu_memory_env(std::vector node_ids); + +namespace cpu_utils { +void eagle_prepare_inputs_padded_kernel_impl( + const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& valid_sampled_tokens_count, + const torch::Tensor& query_start_loc_gpu, + torch::Tensor& token_indices_to_sample, + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs); +void eagle_prepare_next_token_padded_kernel_impl( + const torch::Tensor& sampled_token_ids, + const torch::Tensor& discard_request_mask, + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs); +void eagle_step_slot_mapping_metadata_kernel_impl( + const torch::Tensor& positions, const torch::Tensor& block_table, + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, + torch::Tensor& out_slot_mapping, const int64_t block_size, + const int64_t max_model_len, const int64_t PAD_ID); +void copy_and_expand_eagle_inputs_kernel_impl( + const torch::Tensor& target_token_ids, + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, + torch::Tensor& out_input_ids, torch::Tensor& out_positions, + torch::Tensor& out_is_rejected_token_mask, + torch::Tensor& out_is_masked_token_mask, + torch::Tensor& out_new_token_indices, + torch::Tensor& out_hidden_state_mapping, + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, + const int64_t padding_token_id, const int64_t parallel_drafting_token_id, + const int64_t total_input_tokens, + const int64_t num_padding_slots_per_request, const bool shift_input_ids); +void rejection_greedy_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, + const torch::Tensor& bonus_token_ids, + const std::optional& is_greedy, const int64_t max_spec_len); +void rejection_random_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, + const torch::Tensor& recovered_token_ids, + const torch::Tensor& uniform_probs, + const std::optional& is_greedy, const int64_t max_spec_len, + const int64_t vocab_size, const bool no_draft_probs); +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& cu_num_tokens, + const int64_t replace_from, const int64_t replace_to); +void sample_recovered_tokens_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& inv_q, + const int64_t vocab_size, const bool no_draft_probs); +} // namespace cpu_utils + TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // vLLM custom ops @@ -176,6 +236,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_quick", torch::kCPU, &gelu_quick); +#if (defined(__aarch64__) && !defined(__APPLE__)) + + ops.def( + "activation_lut_bf16(Tensor! out, Tensor input, str activation)" + " -> ()"); + ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16); + +#endif // (defined(__aarch64__) && !defined(__APPLE__)) + // Layernorm // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( @@ -194,7 +263,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "rotary_embedding(Tensor positions, Tensor! query," " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox) -> ()"); + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding); // Quantization @@ -363,6 +433,72 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt " "block_size) -> ()", &compute_slot_mapping_kernel_impl); + + ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); + + // Speculative decoding kernels + ops.def( + "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " + "Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, " + "Tensor(a3!) token_indices_to_sample, " + "Tensor(a4!) num_rejected_tokens_gpu, " + "SymInt num_reqs) -> ()", + &cpu_utils::eagle_prepare_inputs_padded_kernel_impl); + ops.def( + "eagle_prepare_next_token_padded_kernel_impl(" + "Tensor sampled_token_ids, Tensor discard_request_mask, " + "Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, " + "Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, " + "SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()", + &cpu_utils::eagle_prepare_next_token_padded_kernel_impl); + ops.def( + "eagle_step_slot_mapping_metadata_kernel_impl(" + "Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, " + "Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, " + "SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()", + &cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl); + ops.def( + "copy_and_expand_eagle_inputs_kernel_impl(" + "Tensor target_token_ids, Tensor target_positions, " + "Tensor next_token_ids, Tensor(a3!) out_input_ids, " + "Tensor(a4!) out_positions, " + "Tensor(a5!) out_is_rejected_token_mask, " + "Tensor(a6!) out_is_masked_token_mask, " + "Tensor(a7!) out_new_token_indices, " + "Tensor(a8!) out_hidden_state_mapping, " + "Tensor query_start_loc, Tensor query_end_loc, " + "SymInt padding_token_id, SymInt parallel_drafting_token_id, " + "SymInt total_input_tokens, SymInt num_padding_slots_per_request, " + "bool shift_input_ids) -> ()", + &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl); + ops.def( + "rejection_greedy_sample_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor target_argmax, " + "Tensor bonus_token_ids, Tensor? is_greedy, " + "SymInt max_spec_len) -> ()", + &cpu_utils::rejection_greedy_sample_kernel_impl); + ops.def( + "rejection_random_sample_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor? draft_probs, " + "Tensor target_probs, Tensor bonus_token_ids, " + "Tensor recovered_token_ids, Tensor uniform_probs, " + "Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, " + "bool no_draft_probs) -> ()", + &cpu_utils::rejection_random_sample_kernel_impl); + ops.def( + "expand_kernel_impl(Tensor(a0!) output, Tensor input, " + "Tensor cu_num_tokens, SymInt replace_from, " + "SymInt replace_to) -> ()", + &cpu_utils::expand_kernel_impl); + ops.def( + "sample_recovered_tokens_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor? draft_probs, " + "Tensor target_probs, Tensor inv_q, SymInt vocab_size, " + "bool no_draft_probs) -> ()", + &cpu_utils::sample_recovered_tokens_kernel_impl); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/cpu/utils.cpp b/csrc/cpu/utils.cpp index 55d0a3d17572..fd06e01e4d60 100644 --- a/csrc/cpu/utils.cpp +++ b/csrc/cpu/utils.cpp @@ -13,13 +13,80 @@ #include "cpu/utils.hpp" #ifdef VLLM_NUMA_DISABLED -std::string init_cpu_threads_env(const std::string& cpu_ids) { - return std::string( - "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has " - "no effect to setup thread affinity."); -} +void init_cpu_memory_env(std::vector node_ids) {} +#else +void init_cpu_memory_env(std::vector node_ids) { + // Memory node binding + if (numa_available() != -1) { + // Concatenate all node_ids into a single comma-separated string + if (!node_ids.empty()) { + std::string node_ids_str; + for (const int node_id : node_ids) { + if (!node_ids_str.empty()) { + node_ids_str += ","; + } + node_ids_str += std::to_string(node_id); + } -#endif + bitmask* mask = numa_parse_nodestring(node_ids_str.c_str()); + bitmask* src_mask = numa_get_mems_allowed(); + + int pid = getpid(); + + if (mask && src_mask) { + // move all existing pages to the specified numa node. + *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp); + int page_num = numa_migrate_pages(pid, src_mask, mask); + if (page_num == -1) { + TORCH_WARN("numa_migrate_pages failed. errno: " + + std::to_string(errno)); + } + + // Restrict memory allocation to the selected NUMA node(s). + // Enhances memory locality for the threads bound to those NUMA CPUs. + if (node_ids.size() > 1) { + errno = 0; + numa_set_interleave_mask(mask); + if (errno != 0) { + TORCH_WARN("numa_set_interleave_mask failed. errno: " + + std::to_string(errno)); + } else { + TORCH_WARN( + "NUMA binding: Using INTERLEAVE policy for memory " + "allocation across multiple NUMA nodes (nodes: " + + node_ids_str + + "). Memory allocations will be " + "interleaved across the specified NUMA nodes."); + } + } else { + errno = 0; + numa_set_membind(mask); + if (errno != 0) { + TORCH_WARN("numa_set_membind failed. errno: " + + std::to_string(errno)); + } else { + TORCH_WARN( + "NUMA binding: Using MEMBIND policy for memory " + "allocation on the NUMA nodes (" + + node_ids_str + + "). Memory allocations will be " + "strictly bound to these NUMA nodes."); + } + } + + numa_set_strict(1); + + numa_free_nodemask(mask); + numa_free_nodemask(src_mask); + } else { + TORCH_WARN( + "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " + + std::to_string(errno)); + } + } + } +} +#endif // VLLM_NUMA_DISABLED namespace cpu_utils { ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) { diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp index f237bba088b1..2c9e01c60f93 100644 --- a/csrc/cpu/utils.hpp +++ b/csrc/cpu/utils.hpp @@ -54,12 +54,34 @@ struct Counter { }; inline int64_t get_available_l2_size() { +#if defined(__s390x__) + static int64_t size = []() { + uint32_t l2_cache_size = 0; + auto caps = at::cpu::get_cpu_capabilities(); + auto it = caps.find("l2_cache_size"); + if (it != caps.end()) { + l2_cache_size = static_cast(it->second.toInt()); + } + if (l2_cache_size == 0) { + long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE); + if (sys_l2 > 0) { + l2_cache_size = static_cast(sys_l2); + } + } + if (l2_cache_size == 0) { + l2_cache_size = 256 * 1024; + } + return static_cast(l2_cache_size) >> 1; // use 50% of L2 cache + }(); + return size; +#else static int64_t size = []() { auto caps = at::cpu::get_cpu_capabilities(); const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt(); return l2_cache_size >> 1; // use 50% of L2 cache }(); return size; +#endif } template diff --git a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp index 8aa99b3e03aa..0a3c9e9cc7f4 100644 --- a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp @@ -389,20 +389,28 @@ struct Sm90ColOrScalarBroadcastArray { CUTLASS_DEVICE void begin() { - cute::Tensor pred = make_tensor(shape(tCgCol)); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = get<0>(tCcCol(i)) < m; - } - if (!params.col_broadcast) { fill(tCrCol, *(params.ptr_col_array[group])); return; } - // Filter so we don't issue redundant copies over stride-0 modes - // (only works if 0-strides are in same location, which is by construction) - copy_if(pred, filter(tCgCol), filter(tCrCol)); + // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and + // EPI_N are stride-0 for the column broadcast. Slice those modes at + // index 0 to avoid redundant copies AND ensure pred/data consistency + static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast"); + static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast"); + + auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + + cute::Tensor pred = make_tensor(shape(tCgCol_s)); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = get<0>(tCcCol_s(i)) < m; + } + + copy_if(pred, tCgCol_s, tCrCol_s); } template diff --git a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp index 8203d8930ce1..29e6ec41e2a4 100644 --- a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp @@ -382,20 +382,28 @@ struct Sm90ColOrScalarBroadcast { CUTLASS_DEVICE void begin() { - cute::Tensor pred = make_tensor(shape(tCgCol)); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = get<0>(tCcCol(i)) < m; - } - if (!params.col_broadcast) { fill(tCrCol, *(params.ptr_col)); return; } - // Filter so we don't issue redundant copies over stride-0 modes - // (only works if 0-strides are in same location, which is by construction) - copy_if(pred, filter(tCgCol), filter(tCrCol)); + // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and + // EPI_N are stride-0 for the column broadcast. Slice those modes at + // index 0 to avoid redundant copies AND ensure pred/data consistency + static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast"); + static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast"); + + auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + + cute::Tensor pred = make_tensor(shape(tCgCol_s)); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = get<0>(tCcCol_s(i)) < m; + } + + copy_if(pred, tCgCol_s, tCrCol_s); } template diff --git a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu new file mode 100644 index 000000000000..e96017d86dad --- /dev/null +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -0,0 +1,477 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Horizontally-fused DeepseekV4-MLA kernel: + * - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM + * - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged + * cache insert + * + * Structured after `applyMLARopeAndAssignQKVKernelGeneration` in + * TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot + * dispatch choosing Q vs KV work per warp. The per-warp RMSNorm/RoPE + * skeleton is adapted from vllm-deepseek_v4's existing + * `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu). + * + * Assumptions (hard-coded for DeepseekV4 attention): + * HEAD_DIM = 512 + * ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM)) + * NOPE_DIM = 448 + * QUANT_BLOCK = 64 (UE8M0 FP8 quant block) + * FP8_MAX = 448.0f + * is_neox=false (GPT-J interleaved pairs) + * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin + * second along last dim; each half is rope_dim/2 = 32 values) + * + * Cache layout per paged-cache block (block_size tokens): + * [0, bs*576): token data, 448 fp8 + 128 bf16 each + * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + */ + +#include +#include +#include +#include + +#include +#include +#include + +#include "cuda_compat.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" + +#ifndef FINAL_MASK + #define FINAL_MASK 0xffffffffu +#endif + +namespace vllm { +namespace deepseek_v4_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = at::cuda::getCurrentDeviceProperties(); + return props->major * 10 + props->minor; +} +} // namespace + +// ──────────────────────────────────────────────────────────────────────────── +// Constants +// ──────────────────────────────────────────────────────────────────────────── +constexpr int kHeadDim = 512; +constexpr int kRopeDim = 64; +constexpr int kNopeDim = kHeadDim - kRopeDim; // 448 +constexpr int kQuantBlock = 64; +constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7 +constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad) +constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 +constexpr float kFp8Max = 448.0f; + +// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM. +constexpr int kNumLanes = 32; +constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16 + +// ──────────────────────────────────────────────────────────────────────────── +// Small inline helpers +// ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warp4MaxAbs(float val) { + // Reduce absolute max across 4 consecutive lanes (lane id & 3 group). + float peer = __shfl_xor_sync(FINAL_MASK, val, 1); + val = fmaxf(val, peer); + peer = __shfl_xor_sync(FINAL_MASK, val, 2); + val = fmaxf(val, peer); + return val; +} + +template +__device__ __forceinline__ float warpSum(float val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val += __shfl_xor_sync(FINAL_MASK, val, mask, 32); + } + return val; +} + +// ──────────────────────────────────────────────────────────────────────────── +// Kernel +// ──────────────────────────────────────────────────────────────────────────── +// +// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / +// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each +// warp handles one (token, head_slot) pair. head_slot < num_heads_q → +// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV +// branch (RoPE + UE8M0 quant + insert) +// +// With DP padding, q/kv/position_ids can have more rows than slot_mapping. +// The Q branch covers all `num_tokens_full` rows (downstream attention uses +// them). The KV branch only inserts the first `num_tokens_insert` tokens +// (= slot_mapping length) into the paged cache. +// +template +__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( + scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place + scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16 + uint8_t* __restrict__ k_cache, // [num_blocks, block_stride] + int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64 + int64_t const* __restrict__ position_ids, // [N] i64 + float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32 + float const eps, + int const num_tokens_full, // = q.size(0) = kv.size(0) + int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full + int const num_heads_q, // H + int const cache_block_size, // tokens per paged-cache block + int const kv_block_stride) { // bytes per paged-cache block +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + // BF16 _typeConvert specialization is unavailable on pre-Ampere. The + // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a + // no-op stub for sm_70/sm_75 to keep multi-arch builds happy. + if constexpr (std::is_same_v) { + return; + } else { +#endif + using Converter = vllm::_typeConvert; + + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; + + int const total_slots_per_token = num_heads_q + 1; + int const tokenIdx = globalWarpIdx / total_slots_per_token; + int const slotIdx = globalWarpIdx % total_slots_per_token; + if (tokenIdx >= num_tokens_full) return; + + bool const isKV = (slotIdx == num_heads_q); + // KV branch: skip DP-padded tokens (no slot reserved for them). + if (isKV && tokenIdx >= num_tokens_insert) return; + + // PDL: wait for predecessor kernel (upstream q/kv producer) to signal + // before touching any global memory. No-op when PDL is not enabled on + // the launch. The CUDA runtime wrapper emits the griddepcontrol.wait + // PTX with the required memory clobber internally. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // Dim range this lane owns within the 512-wide head. + int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16 + + // ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ──── + float elements[kElemsPerLane]; + float sumOfSquares = 0.0f; + + scalar_t_in const* src_ptr; + if (isKV) { + src_ptr = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; + } else { + int64_t const q_row_offset = + (static_cast(tokenIdx) * num_heads_q + slotIdx) * kHeadDim + + dim_base; + src_ptr = q_inout + q_row_offset; + } + + // Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector + // type and bitcast to scalar_t_in packed pairs for conversion. + uint4 v0 = *reinterpret_cast(src_ptr); + uint4 v1 = *reinterpret_cast(src_ptr + 8); + + { + typename Converter::packed_hip_type const* p0 = + reinterpret_cast(&v0); + typename Converter::packed_hip_type const* p1 = + reinterpret_cast(&v1); +// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4. +#pragma unroll + for (int i = 0; i < 4; i++) { + float2 f2 = Converter::convert(p0[i]); + elements[2 * i] = f2.x; + elements[2 * i + 1] = f2.y; + } +#pragma unroll + for (int i = 0; i < 4; i++) { + float2 f2 = Converter::convert(p1[i]); + elements[8 + 2 * i] = f2.x; + elements[8 + 2 * i + 1] = f2.y; + } + } + + // ── Q branch: RMSNorm with no weight (has_weight=False) ───────────────── + // Variance + rsqrt + multiply all in fp32, no intermediate bf16 round. + // The downstream bf16 round only happens at the final store. + if (!isKV) { +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + sumOfSquares += elements[i] * elements[i]; + } + sumOfSquares = warpSum(sumOfSquares); + float const rms_rcp = + rsqrtf(sumOfSquares / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + elements[i] = elements[i] * rms_rcp; + } + } + + // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────── + // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage). + bool const is_rope_lane = dim_base >= kNopeDim; + if (is_rope_lane) { + int64_t const pos = position_ids[tokenIdx]; + constexpr int kHalfRope = kRopeDim / 2; // 32 + float const* cos_ptr = cos_sin_cache + pos * kRopeDim; + float const* sin_ptr = cos_ptr + kHalfRope; + + int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16 +#pragma unroll + for (int p = 0; p < kElemsPerLane / 2; p++) { + int const pair_dim = rope_local_base + 2 * p; + int const half_idx = pair_dim / 2; + float const cos_v = VLLM_LDG(cos_ptr + half_idx); + float const sin_v = VLLM_LDG(sin_ptr + half_idx); + float const x_even = elements[2 * p]; + float const x_odd = elements[2 * p + 1]; + elements[2 * p] = x_even * cos_v - x_odd * sin_v; + elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v; + } + } + + // ═══════════════════════════════════════════════════════════════════════ + // Q branch: cast to bf16 and store back in place. + // ═══════════════════════════════════════════════════════════════════════ + if (!isKV) { + uint4 out0, out1; + typename Converter::packed_hip_type* po0 = + reinterpret_cast(&out0); + typename Converter::packed_hip_type* po1 = + reinterpret_cast(&out1); +#pragma unroll + for (int i = 0; i < 4; i++) { + po0[i] = Converter::convert( + make_float2(elements[2 * i], elements[2 * i + 1])); + } +#pragma unroll + for (int i = 0; i < 4; i++) { + po1[i] = Converter::convert( + make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1])); + } + scalar_t_in* dst = + q_inout + + (static_cast(tokenIdx) * num_heads_q + slotIdx) * kHeadDim + + dim_base; + *reinterpret_cast(dst) = out0; + *reinterpret_cast(dst + 8) = out1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + + // ═══════════════════════════════════════════════════════════════════════ + // KV branch. + // ═══════════════════════════════════════════════════════════════════════ + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id < 0) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + + int64_t const block_idx = slot_id / cache_block_size; + int64_t const pos_in_block = slot_id % cache_block_size; + uint8_t* block_base = + k_cache + block_idx * static_cast(kv_block_stride); + uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; + uint8_t* token_scale_ptr = + block_base + static_cast(cache_block_size) * kTokenDataBytes + + pos_in_block * kScaleBytesPerToken; + + // Round K to bf16 first, matching the unfused reference path where K is + // materialized as bf16 before K quantization. absmax, clamp, and FP8 + // quant below all run on these bf16-rounded values. +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + elements[i] = Converter::convert(Converter::convert(elements[i])); + } + + // Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective + // shuffle requires full participation). RoPE lanes contribute garbage, + // but their values are gated out below via `!is_rope_lane`. + float local_absmax = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + local_absmax = fmaxf(local_absmax, fabsf(elements[i])); + } + float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f); + float const exponent = ceilf(log2f(absmax / kFp8Max)); + float const inv_scale = exp2f(-exponent); + + if (!is_rope_lane) { + // ── NoPE lane: UE8M0 FP8 quant ─────────────────────────────────────── + uint8_t out_bytes[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + float scaled = elements[i] * inv_scale; + scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max); + __nv_fp8_storage_t s = + __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3); + out_bytes[i] = static_cast(s); + } + // One 16-byte STG per lane. + *reinterpret_cast(token_fp8_ptr + dim_base) = + *reinterpret_cast(out_bytes); + + // Lane (4k) of each 4-lane group writes the scale byte for block k<7. + if ((laneId & 3) == 0) { + int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes + float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); + token_scale_ptr[q_block_idx] = static_cast(encoded); + } + // Lane 0 also writes the padding byte at index 7. + if (laneId == 0) { + token_scale_ptr[kNumQuantBlocks] = 0; // pad + } + } else { + // ── RoPE lane: cast back to bf16 and store to cache bf16 tail ──────── + uint4 out0, out1; + typename Converter::packed_hip_type* po0 = + reinterpret_cast(&out0); + typename Converter::packed_hip_type* po1 = + reinterpret_cast(&out1); +#pragma unroll + for (int i = 0; i < 4; i++) { + po0[i] = Converter::convert( + make_float2(elements[2 * i], elements[2 * i + 1])); + } +#pragma unroll + for (int i = 0; i < 4; i++) { + po1[i] = Converter::convert( + make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1])); + } + int const rope_local_base = dim_base - kNopeDim; // in [0, 64) + scalar_t_in* bf16_dst = + reinterpret_cast(token_bf16_ptr) + rope_local_base; + *reinterpret_cast(bf16_dst) = out0; + *reinterpret_cast(bf16_dst + 8) = out1; + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Launch wrapper +// ──────────────────────────────────────────────────────────────────────────── +template +void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( + scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache, + int64_t const* slot_mapping, int64_t const* position_ids, + float const* cos_sin_cache, float const eps, int const num_tokens_full, + int const num_tokens_insert, int const num_heads_q, + int const cache_block_size, int const kv_block_stride, + cudaStream_t stream) { + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int64_t const total_warps = + static_cast(num_tokens_full) * (num_heads_q + 1); + int const grid = + static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + + // PDL: enable programmatic stream serialization whenever the hardware + // supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable, + // so leave numAttrs = 0 and launch as a regular kernel. + static int const sm_version = getSMVersion(); + // Host-side guard: the device kernel body is compiled as a no-op for + // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert is + // unavailable there. Refuse the launch loudly instead of silently + // skipping the work. + TORCH_CHECK( + sm_version >= 80, + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ " + "(Ampere or newer); got sm_", + sm_version); + cudaLaunchConfig_t config; + config.gridDim = dim3(grid); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + + cudaLaunchKernelEx( + &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel, + q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, + num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, + kv_block_stride); +} + +} // namespace deepseek_v4_fused_ops +} // namespace vllm + +// ──────────────────────────────────────────────────────────────────────────── +// Torch op wrapper +// ──────────────────────────────────────────────────────────────────────────── +void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::Tensor& q, // [N, H, 512] bf16, in place + torch::Tensor const& kv, // [N, 512] bf16 (read-only) + torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8 + torch::Tensor const& slot_mapping, // [N] int64 + torch::Tensor const& position_ids, // [N] int64 + torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 + double eps, int64_t cache_block_size) { + TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA"); + TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA"); + TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA"); + TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64, + "slot_mapping must be int64 CUDA"); + TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64, + "position_ids must be int64 CUDA"); + TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA"); + TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]"); + TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); + TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match"); + TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8"); + TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, + "cos_sin_cache shape [max_pos, 64]"); + TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32, + "cos_sin_cache must be float32"); + + // With DP padding, slot_mapping can be shorter than q/kv/positions. + // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them); + // KV quant+insert runs only on the first slot_mapping.size(0) rows. + int const num_tokens_full = static_cast(q.size(0)); + int const num_tokens_insert = static_cast(slot_mapping.size(0)); + TORCH_CHECK(static_cast(kv.size(0)) == num_tokens_full && + static_cast(position_ids.size(0)) == num_tokens_full, + "q/kv/position_ids row counts must match"); + TORCH_CHECK(num_tokens_insert <= num_tokens_full, + "slot_mapping must not exceed q row count"); + int const num_heads_q = static_cast(q.size(1)); + int const cache_block_size_i = static_cast(cache_block_size); + int const kv_block_stride = static_cast(k_cache.stride(0)); + + at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + auto stream = at::cuda::getCurrentCUDAStream(); + + VLLM_DISPATCH_HALF_TYPES( + q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { + using qkv_scalar_t = scalar_t; + vllm::deepseek_v4_fused_ops:: + launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(kv.data_ptr()), + reinterpret_cast(k_cache.data_ptr()), + reinterpret_cast(slot_mapping.data_ptr()), + reinterpret_cast(position_ids.data_ptr()), + cos_sin_cache.data_ptr(), static_cast(eps), + num_tokens_full, num_tokens_insert, num_heads_q, + cache_block_size_i, kv_block_stride, stream); + }); +} diff --git a/csrc/fused_qknorm_rope_kernel.cu b/csrc/fused_qknorm_rope_kernel.cu index a51e1a347e1d..0bf48fd3e831 100644 --- a/csrc/fused_qknorm_rope_kernel.cu +++ b/csrc/fused_qknorm_rope_kernel.cu @@ -19,8 +19,10 @@ #include #include +#include #include +#include "async_util.cuh" #include "cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" @@ -86,6 +88,9 @@ inline __device__ __host__ T divUp(T m, T n) { } // namespace tensorrt_llm::common namespace tensorrt_llm::kernels { + +using namespace vllm::cuda_async; + // NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation, // with added support for passing the cos_sin_cache as an input. // https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -301,6 +306,237 @@ __global__ void fusedQKNormRopeKernel( #endif } +// Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for +// the same token, sharing cos/sin from shared memory via cp.async. +// When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads, +// hiding global-memory latency and improving occupancy for large batches. +template +__global__ void fusedQKNormRopeKernelNTokenHeads( + void* qkv_void, int const num_heads_q, int const num_heads_k, + int const num_heads_v, float const eps, void const* q_weight_void, + void const* k_weight_void, void const* cos_sin_cache_void, + int64_t const* position_ids, int const num_tokens, int const rotary_dim) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + if constexpr ((std::is_same_v) || + std::is_same_v) { + return; + } else { +#endif + + using Converter = vllm::_typeConvert; + static_assert(Converter::exists, + "Input QKV data type is not supported for this CUDA " + "architecture or toolkit version."); + using T_in = typename Converter::hip_type; + using T2_in = typename Converter::packed_hip_type; + + using CacheConverter = vllm::_typeConvert; + static_assert(CacheConverter::exists, + "Cache data type is not supported for this CUDA architecture " + "or toolkit version."); + using T_cache = typename CacheConverter::hip_type; + + extern __shared__ char smem_storage[]; + // Shared memory layout: + // [0, cos_sin_bytes) : cos/sin for each warp (warpsPerBlock * + // rotary_dim * sizeof(T_cache)) + // [cos_sin_bytes, ...) : QKV tiles + // per warp (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes) + T_cache* const smem = reinterpret_cast(smem_storage); + + T_in* qkv = reinterpret_cast(qkv_void); + T_in const* q_weight = reinterpret_cast(q_weight_void); + T_in const* k_weight = reinterpret_cast(k_weight_void); + T_cache const* cos_sin_cache = + reinterpret_cast(cos_sin_cache_void); + + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + + int const total_qk_heads = num_heads_q + num_heads_k; + int const num_heads = num_heads_q + num_heads_k + num_heads_v; + int const head_chunks_per_token = + (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP; + + int const warp_global = blockIdx.x * warpsPerBlock + warpId; + int const tokenIdx = warp_global / head_chunks_per_token; + int const headChunk = warp_global % head_chunks_per_token; + int const first_head = headChunk * HEADS_PER_WARP; + int const num_heads_this_warp = + (first_head + HEADS_PER_WARP <= total_qk_heads) + ? HEADS_PER_WARP + : (total_qk_heads - first_head); + + if (tokenIdx >= num_tokens) return; + + static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64"); + constexpr int numElemsPerThread = head_dim / 32; + constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16); + static_assert(elemSizeBytes % 4 == 0, + "elemSizeBytes must be a multiple of 4"); + constexpr int vecSize = elemSizeBytes / 4; + using vec_T = typename tensorrt_llm::common::packed_as::type; + + int const cos_sin_bytes = + warpsPerBlock * rotary_dim * static_cast(sizeof(T_cache)); + int const qkv_tile_bytes = 32 * elemSizeBytes; + char* const this_warp_head_smem = + smem_storage + cos_sin_bytes + + warpId * (HEADS_PER_WARP * qkv_tile_bytes); + + // === Group 0: async load all heads' QKV into smem (issued first). === + for (int k = 0; k < num_heads_this_warp; ++k) { + int const localHeadIdx = first_head + k; + bool const isQ = localHeadIdx < num_heads_q; + int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; + int offWarp; + if (isQ) { + offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; + } else { + offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + + headIdx * head_dim; + } + int const offThread = offWarp + laneId * numElemsPerThread; + char* smem_dst = + this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; + cp_async_shared_global_ca(smem_dst, + reinterpret_cast(&qkv[offThread]), + elemSizeBytes); + } + cp_async_commit_group(); // commit group 0 (QKV) + + // === Group 1: async load cos/sin into smem (issued second). === + int64_t const pos_id = position_ids[tokenIdx]; + T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim; + int const copy_bytes = rotary_dim * static_cast(sizeof(T_cache)); + int const num_copies = (copy_bytes + 15) / 16; + for (int copyId = laneId; copyId < num_copies; copyId += 32) { + char* smem_ptr = + reinterpret_cast(&smem[warpId * rotary_dim]) + copyId * 16; + const char* glob_ptr = + reinterpret_cast(cache_ptr) + copyId * 16; + cp_async_shared_global_16_cg(smem_ptr, glob_ptr); + } + cp_async_commit_group(); // commit group 1 (cos/sin) + + // wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done. + cp_async_wait_group<1>(); + + float elements[numElemsPerThread]; + float elements2[numElemsPerThread]; + int const rotary_lanes = rotary_dim / numElemsPerThread; + int const embed_dim = rotary_dim / 2; + T_cache const* const cos_smem = &smem[warpId * rotary_dim]; + T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim]; + + // Preload weights into registers once, reused across all heads. + float q_w[numElemsPerThread]; + float k_w[numElemsPerThread]; +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + int const dim = laneId * numElemsPerThread + i; + q_w[i] = Converter::convert(q_weight[dim]); + k_w[i] = Converter::convert(k_weight[dim]); + } + + for (int k = 0; k < num_heads_this_warp; ++k) { + int const localHeadIdx = first_head + k; + bool const isQ = localHeadIdx < num_heads_q; + int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; + + int offsetWarp; + if (isQ) { + offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; + } else { + offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + + headIdx * head_dim; + } + int const offsetThread = offsetWarp + laneId * numElemsPerThread; + + // === Part 1: QK Norm (read from smem; group 0 already done). === + float sumOfSquares = 0.0f; + { + char const* smem_src = + this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; + vec_T vec = *reinterpret_cast(smem_src); + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); +#pragma unroll + for (int i = 0; i < num_packed_elems; i++) { + T2_in packed_val = *(reinterpret_cast(&vec) + i); + float2 vals = Converter::convert(packed_val); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + + sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(head_dim) + eps); + +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]); + } + + // On first head: wait for group 1 (cos/sin) before RoPE. + if (k == 0) cp_async_wait_group<0>(); + + // === Part 2: RoPE using cos/sin from shared memory. === + if (laneId < rotary_lanes) { + if constexpr (interleave) { +#pragma unroll + for (int i = 0; i < numElemsPerThread / 2; ++i) { + int const idx0 = 2 * i; + int const idx1 = 2 * i + 1; + int const dim_idx = laneId * numElemsPerThread + idx0; + float const val0 = elements[idx0]; + float const val1 = elements[idx1]; + int const half_dim = dim_idx / 2; + float const cos_val = CacheConverter::convert(cos_smem[half_dim]); + float const sin_val = CacheConverter::convert(sin_smem[half_dim]); + elements[idx0] = val0 * cos_val - val1 * sin_val; + elements[idx1] = val0 * sin_val + val1 * cos_val; + } + } else { + __syncwarp(); + int const pairOffset = (rotary_dim / 2) / numElemsPerThread; +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset); + if (laneId < pairOffset) elements2[i] = -elements2[i]; + int dim_idx = laneId * numElemsPerThread + i; + dim_idx = (dim_idx * 2) % rotary_dim; + int const half_dim = dim_idx / 2; + float const cos_val = CacheConverter::convert(cos_smem[half_dim]); + float const sin_val = CacheConverter::convert(sin_smem[half_dim]); + elements[i] = elements[i] * cos_val + elements2[i] * sin_val; + } + __syncwarp(); + } + } + + // Store. + { + vec_T vec; + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); +#pragma unroll + for (int i = 0; i < num_packed_elems; i++) { + T2_in packed_val = Converter::convert( + make_float2(elements[2 * i], elements[2 * i + 1])); + *(reinterpret_cast(&vec) + i) = packed_val; + } + *reinterpret_cast(&qkv[offsetThread]) = vec; + } + } + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif +} + // Borrowed from // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568 #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ @@ -321,15 +557,12 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, void const* cos_sin_cache, bool const interleave, int64_t const* position_ids, cudaStream_t stream) { constexpr int blockSize = 256; - int const warpsPerBlock = blockSize / 32; int const totalQKHeads = num_heads_q + num_heads_k; int const totalWarps = num_tokens * totalQKHeads; - int const gridSize = common::divUp(totalWarps, warpsPerBlock); dim3 gridDim(gridSize); dim3 blockDim(blockSize); - switch (head_dim) { case 64: DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { @@ -360,6 +593,118 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, "Unsupported head dimension for fusedQKNormRope: ", head_dim); } } + +// Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8). +// When token_heads_per_warp == 1, delegates to the 1-head baseline above. +template +void launchFusedQKNormRopeNTokenHeads( + void* qkv, int const num_tokens, int const num_heads_q, + int const num_heads_k, int const num_heads_v, int const head_dim, + int const rotary_dim, float const eps, void const* q_weight, + void const* k_weight, void const* cos_sin_cache, bool const interleave, + int64_t const* position_ids, int const token_heads_per_warp, + cudaStream_t stream) { + TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 || + token_heads_per_warp == 4 || token_heads_per_warp == 8, + "token_heads_per_warp must be 1, 2, 4, or 8, got ", + token_heads_per_warp); + + // token_heads_per_warp == 1: delegate to the 1-head baseline kernel. + if (token_heads_per_warp == 1) { + launchFusedQKNormRope( + qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, + rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, + position_ids, stream); + return; + } + + // NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks. + // If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last + // cp.async would write past the shared memory allocation. + // Fall back to the base kernel instead of failing. + { + size_t const rotary_bytes = + static_cast(rotary_dim) * + (std::is_same_v ? sizeof(float) : 2u); + if (rotary_bytes % 16 != 0) { + launchFusedQKNormRope( + qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, + rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, + position_ids, stream); + return; + } + } + + constexpr int blockSize = 256; + int const warpsPerBlock = blockSize / 32; + int const totalQKHeads = num_heads_q + num_heads_k; + // Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem. + int const head_chunks_per_token = + (totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp; + int const total_warps = num_tokens * head_chunks_per_token; + int const gridSize = common::divUp(total_warps, warpsPerBlock); + dim3 gridDim(gridSize); + dim3 blockDim(blockSize); + // Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same + // layout). + size_t const cache_elem_size = + std::is_same_v ? sizeof(float) : 2u; + // QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2) + // = 2*head_dim bytes. + size_t const qkv_smem_per_warp = static_cast(token_heads_per_warp) * + 2u * static_cast(head_dim); + size_t const smem_bytes = + warpsPerBlock * static_cast(rotary_dim) * cache_elem_size + + warpsPerBlock * qkv_smem_per_warp; + +#define LAUNCH_N_TOKEN_HEADS(N) \ + do { \ + switch (head_dim) { \ + case 64: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + case 128: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + case 256: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \ + } \ + } while (0) + + if (token_heads_per_warp == 2) { + LAUNCH_N_TOKEN_HEADS(2); + } else if (token_heads_per_warp == 4) { + LAUNCH_N_TOKEN_HEADS(4); + } else if (token_heads_per_warp == 8) { + LAUNCH_N_TOKEN_HEADS(8); + } +#undef LAUNCH_N_TOKEN_HEADS +} + } // namespace tensorrt_llm::kernels void fused_qk_norm_rope( @@ -374,7 +719,8 @@ void fused_qk_norm_rope( torch::Tensor& k_weight, // RMSNorm weights for key [head_dim] torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim] bool is_neox, // Whether RoPE is applied in Neox style - torch::Tensor& position_ids // Position IDs for RoPE [num_tokens] + torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens] + int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value ) { // Input validation CHECK_INPUT(qkv); @@ -414,15 +760,48 @@ void fused_qk_norm_rope( qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total number of heads and head dimension"); - auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + auto device_id = qkv.get_device(); + auto stream = at::cuda::getCurrentCUDAStream(device_id); + + // Select token_heads_per_warp: forced value if >0, else auto-select. + // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures, + // fall back to token_heads_per_warp=1 (base kernel) until profiled. + int token_heads_per_warp; + if (forced_token_heads_per_warp > 0) { // only support SM80+ + token_heads_per_warp = static_cast(forced_token_heads_per_warp); + } else { + token_heads_per_warp = 1; + auto* dev_prop = at::cuda::getDeviceProperties(device_id); + int sm_version = dev_prop->major * 10 + dev_prop->minor; + int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k); + if (sm_version == 90) { + if (head_dim >= 256) { + if (total_qk_units < 4096LL) { + token_heads_per_warp = 1; + } else if (total_qk_units < 8192LL) { + token_heads_per_warp = 2; + } else { + token_heads_per_warp = 4; + } + } else { + if (total_qk_units < 10240LL) { + token_heads_per_warp = 1; + } else if (total_qk_units < 40960LL) { + token_heads_per_warp = 4; + } else { + token_heads_per_warp = 8; + } + } + } + } VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] { using qkv_scalar_t = scalar_t; VLLM_DISPATCH_FLOATING_TYPES( cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] { using cache_scalar_t = scalar_t; - tensorrt_llm::kernels::launchFusedQKNormRope( + tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads< + qkv_scalar_t, cache_scalar_t>( qkv.data_ptr(), static_cast(num_tokens), static_cast(num_heads_q), static_cast(num_heads_k), static_cast(num_heads_v), static_cast(head_dim), @@ -430,7 +809,7 @@ void fused_qk_norm_rope( q_weight.data_ptr(), k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox, reinterpret_cast(position_ids.data_ptr()), - stream); + token_heads_per_warp, stream); }); }); -} \ No newline at end of file +} diff --git a/csrc/layernorm_kernels.cu b/csrc/layernorm_kernels.cu index 9766103f7646..e617e45dc58b 100644 --- a/csrc/layernorm_kernels.cu +++ b/csrc/layernorm_kernels.cu @@ -77,7 +77,8 @@ __global__ void rms_norm_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j]; + float w = static_cast(src2.val[j]); + dst.val[j] = static_cast(x * s_variance * w); } v_out[i] = dst; } @@ -134,10 +135,17 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { int id = blockIdx.x * vec_hidden_size + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; - _f16Vec temp = residual_v[id]; - temp *= s_variance; - temp *= weight_v[idx]; - input_v[strided_id] = temp; + _f16Vec res = residual_v[id]; + _f16Vec w = weight_v[idx]; + _f16Vec out; + using Converter = _typeConvert; +#pragma unroll + for (int j = 0; j < width; ++j) { + float x = Converter::convert(res.data[j]); + float wf = Converter::convert(w.data[j]); + out.data[j] = Converter::convert(x * s_variance * wf); + } + input_v[strided_id] = out; } } @@ -174,8 +182,8 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - input[blockIdx.x * input_stride + idx] = - ((scalar_t)(x * s_variance)) * weight[idx]; + float w = (float)weight[idx]; + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); } } diff --git a/csrc/layernorm_quant_kernels.cu b/csrc/layernorm_quant_kernels.cu index f96386252c31..b2f546b3db1d 100644 --- a/csrc/layernorm_quant_kernels.cu +++ b/csrc/layernorm_quant_kernels.cu @@ -65,9 +65,16 @@ __global__ void rms_norm_static_fp8_quant_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j]; + float w = static_cast(src2.val[j]); + // Round normalized result through scalar_t to match the precision of the + // unfused composite (rms_norm writes scalar_t, then + // static_scaled_fp8_quant re-loads it as float before FP8 conversion). + // Without this round, the fused path is strictly more accurate and + // disagrees with the composite at exact E4M3 quantization tie boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] = - scaled_fp8_conversion(out_norm, scale_inv); + scaled_fp8_conversion(static_cast(out_norm), + scale_inv); } } } @@ -127,13 +134,21 @@ fused_add_rms_norm_static_fp8_quant_kernel( for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { int id = blockIdx.x * vec_hidden_size + idx; - _f16Vec temp = residual_v[id]; - temp *= s_variance; - temp *= weight_v[idx]; + _f16Vec res = residual_v[id]; + _f16Vec w = weight_v[idx]; + using Converter = _typeConvert; + using HipT = typename Converter::hip_type; #pragma unroll for (int i = 0; i < width; ++i) { - out[id * width + i] = - scaled_fp8_conversion(float(temp.data[i]), scale_inv); + float x = Converter::convert(res.data[i]); + float wf = Converter::convert(w.data[i]); + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. We use the + // backend's hip_type for the intermediate since c10::Half/BFloat16 has + // ambiguous conversions on CUDA and no implicit conversion on ROCm. + HipT out_norm_h = Converter::convert(x * s_variance * wf); + out[id * width + i] = scaled_fp8_conversion( + Converter::convert(out_norm_h), scale_inv); } } } @@ -176,9 +191,12 @@ fused_add_rms_norm_static_fp8_quant_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx]; - out[blockIdx.x * hidden_size + idx] = - scaled_fp8_conversion(out_norm, scale_inv); + float w = (float)weight[idx]; + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); + out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion( + static_cast(out_norm), scale_inv); } } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 8153102c598a..176cd5006335 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -134,4 +134,13 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, torch::stable::Tensor& input, torch::stable::Tensor& input_global_scale); +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + #endif diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu new file mode 100644 index 000000000000..8a493fdf22c3 --- /dev/null +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -0,0 +1,468 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100. + * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups. + */ + +#include +#include +#include "libtorch_stable/torch_utils.h" + +#include + +#include "cutlass_extensions/common.hpp" + +#include "cute/tensor.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" + +#include "cutlass/util/packed_stride.hpp" +#include + +using namespace cute; + +// Offset-computation kernel for MXFP4 grouped GEMM (group size 32). +template +__global__ void __mxfp4_get_group_gemm_starts( + ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets, + ElementSF** a_scales_offsets, ElementSF** b_scales_offsets, + LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int, + ElementAB* a_base_as_int, ElementAB* b_base_as_int, + ElementC* out_base_as_int, ElementSF* a_scales_base_as_int, + ElementSF* b_scales_base_as_int, const int32_t* expert_offsets, + const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes, + int64_t* a_strides, int64_t* b_strides, int64_t* c_strides, + const int64_t a_stride_val, const int64_t b_stride_val, + const int64_t c_stride_val, const int K, const int N) { + int64_t expert_id = threadIdx.x; + if (expert_id >= gridDim.x * blockDim.x) { + return; + } + int64_t expert_offset = static_cast(expert_offsets[expert_id]); + int64_t sf_offset = static_cast(sf_offsets[expert_id]); + int64_t group_size = 32; + int64_t m = static_cast(problem_sizes_as_shapes[expert_id * 3]); + int64_t n = static_cast(problem_sizes_as_shapes[expert_id * 3 + 1]); + int64_t k = static_cast(problem_sizes_as_shapes[expert_id * 3 + 2]); + assert((m >= 0 && n == N && k == K && k % 2 == 0) && + "unexpected problem sizes"); + + int64_t half_k = static_cast(k / 2); + int64_t group_k = static_cast(k / group_size); + // Shape of A as uint8/byte = [M, K // 2] + a_offsets[expert_id] = a_base_as_int + expert_offset * half_k; + // Shape of B as uint8/byte = [E, N, K // 2] + b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k; + // Shape of C = [M, N] + out_offsets[expert_id] = out_base_as_int + expert_offset * n; + // Shape of a_scale = [sum(sf_sizes), K // group_size] + a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k; + + assert((reinterpret_cast(a_scales_offsets[expert_id]) % 128) == + 0 && + "TMA requires 128-byte alignment"); + + // Shape of B scale = [E, N, K // group_size] + b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k; + assert((reinterpret_cast(b_scales_offsets[expert_id]) % 128) == + 0 && + "TMA requires 128-byte alignment"); + + // Initialize strides + a_strides[expert_id] = a_stride_val; + b_strides[expert_id] = b_stride_val; + c_strides[expert_id] = c_stride_val; + + LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id; + LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id; + + *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape( + static_cast(m), static_cast(n), static_cast(k), 1)); + *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape( + static_cast(m), static_cast(n), static_cast(k), 1)); +} + +#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \ + TENSOR_C_TYPE, C_TYPE, LayoutSFA, \ + LayoutSFB, ScaleConfig) \ + else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \ + __mxfp4_get_group_gemm_starts \ + <<<1, num_experts, 0, stream>>>( \ + static_cast(a_starts.data_ptr()), \ + static_cast(b_starts.data_ptr()), \ + static_cast(out_starts.data_ptr()), \ + static_cast(a_scales_starts.data_ptr()), \ + static_cast(b_scales_starts.data_ptr()), \ + reinterpret_cast(layout_sfa.data_ptr()), \ + reinterpret_cast(layout_sfb.data_ptr()), \ + static_cast(a_tensors.data_ptr()), \ + static_cast(b_tensors.data_ptr()), \ + static_cast(out_tensors.data_ptr()), \ + static_cast(a_scales.data_ptr()), \ + static_cast(b_scales.data_ptr()), \ + static_cast(expert_offsets.data_ptr()), \ + static_cast(sf_offsets.data_ptr()), \ + static_cast(problem_sizes.data_ptr()), \ + static_cast(a_strides.data_ptr()), \ + static_cast(b_strides.data_ptr()), \ + static_cast(c_strides.data_ptr()), a_stride_val, \ + b_stride_val, c_stride_val, K, N); \ + } + +template +void mxfp4_run_get_group_gemm_starts( + const torch::stable::Tensor& a_starts, + const torch::stable::Tensor& b_starts, + const torch::stable::Tensor& out_starts, + const torch::stable::Tensor& a_scales_starts, + const torch::stable::Tensor& b_scales_starts, + const torch::stable::Tensor& layout_sfa, + const torch::stable::Tensor& layout_sfb, + const torch::stable::Tensor& a_strides, + const torch::stable::Tensor& b_strides, + const torch::stable::Tensor& c_strides, int64_t a_stride_val, + int64_t b_stride_val, int64_t c_stride_val, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& out_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& sf_offsets, + torch::stable::Tensor const& problem_sizes, int M, int N, int K) { + int num_experts = (int)expert_offsets.size(0); + auto stream = get_current_cuda_stream(a_tensors.get_device_index()); + + STD_TORCH_CHECK(out_tensors.size(1) == N, + "Output tensor shape doesn't match expected shape"); + STD_TORCH_CHECK(K / 2 == b_tensors.size(2), + "b_tensors(dim = 2) and a_tensors(dim = 1) trailing" + " dimension must match"); + if (false) { + } + // MXFP4 uses E8M0 (float_ue8m0_t) scale factors + __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, + torch::headeronly::ScalarType::BFloat16, + cutlass::bfloat16_t, LayoutSFA, LayoutSFB, + ScaleConfig) + __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, + torch::headeronly::ScalarType::Half, half, + LayoutSFA, LayoutSFB, ScaleConfig) + else { + STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); + } +} + +template +void run_mxfp4_blockwise_scaled_group_mm_sm100( + torch::stable::Tensor& output, const torch::stable::Tensor& a, + const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + using ProblemShape = + cutlass::gemm::GroupProblemShape>; + using ElementType = cutlass::float_e2m1_t; + using ElementSFType = cutlass::float_ue8m0_t; + using ElementA = cutlass::mx_float4_t; + using ElementB = cutlass::mx_float4_t; + + using ElementC = OutType; + using ElementD = ElementC; + using ElementAccumulator = float; + // Layout definitions + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = LayoutC; + + static constexpr int AlignmentA = 32; + static constexpr int AlignmentB = 32; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + // Architecture definitions + using ArchTag = cutlass::arch::Sm100; + using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; + using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + using StageCountType = cutlass::gemm::collective::StageCountAuto; + + using ClusterShape = Shape<_1, _1, _1>; + struct MMA1SMConfig { + using MmaTileShape = Shape<_128, _128, _128>; + using KernelSchedule = + cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; + }; + + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape, + ClusterShape, Shape<_128, _64>, ElementAccumulator, + ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, + LayoutC*, AlignmentD, + typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA, + ElementB, LayoutB*, AlignmentB, ElementAccumulator, + typename MMA1SMConfig::MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + typename MMA1SMConfig::KernelSchedule>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + + using LayoutSFA = + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LayoutSFB = + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + using ScaleConfig = + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; + int num_experts = static_cast(expert_offsets.size(0)); + + torch::stable::Tensor a_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor out_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor a_scales_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_scales_ptrs = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor layout_sfa = torch::stable::empty( + {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, + a.device()); + torch::stable::Tensor layout_sfb = torch::stable::empty( + {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, + a.device()); + torch::stable::Tensor a_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor b_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + torch::stable::Tensor c_strides1 = + torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, + std::nullopt, a.device()); + + mxfp4_run_get_group_gemm_starts( + a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa, + layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2, + b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale, + b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K); + + // Create an instance of the GEMM + Gemm gemm_op; + + UnderlyingProblemShape* problem_sizes_as_shapes = + static_cast(problem_sizes.data_ptr()); + + // Set the Scheduler info + cutlass::KernelHardwareInfo hw_info; + using RasterOrderOptions = typename cutlass::gemm::kernel::detail:: + PersistentTileSchedulerSm100GroupParams< + typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; + typename Gemm::GemmKernel::TileSchedulerArguments scheduler; + scheduler.raster_order = RasterOrderOptions::AlongM; + hw_info.device_id = a.get_device_index(); + static std::unordered_map cached_sm_counts; + if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { + cached_sm_counts[hw_info.device_id] = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + } + hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX); + + // Mainloop Arguments + typename GemmKernel::MainloopArguments mainloop_args{ + static_cast(a_ptrs.data_ptr()), + static_cast(a_strides1.data_ptr()), + static_cast(b_ptrs.data_ptr()), + static_cast(b_strides1.data_ptr()), + static_cast(a_scales_ptrs.data_ptr()), + reinterpret_cast(layout_sfa.data_ptr()), + static_cast(b_scales_ptrs.data_ptr()), + reinterpret_cast(layout_sfb.data_ptr())}; + + // Epilogue Arguments + typename GemmKernel::EpilogueArguments epilogue_args{ + {}, // epilogue.thread + nullptr, + static_cast(c_strides1.data_ptr()), + static_cast(out_ptrs.data_ptr()), + static_cast(c_strides1.data_ptr())}; + auto& fusion_args = epilogue_args.thread; + // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C + fusion_args.alpha_ptr = nullptr; + fusion_args.beta_ptr = nullptr; + fusion_args.alpha = 1.0f; + fusion_args.alpha_ptr_array = nullptr; + fusion_args.dAlpha = {_0{}, _0{}, 0}; + fusion_args.beta = 0.0f; + fusion_args.beta_ptr_array = nullptr; + fusion_args.dBeta = {_0{}, _0{}, 0}; + + // Gemm Arguments + typename GemmKernel::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {num_experts, problem_sizes_as_shapes, nullptr}, + mainloop_args, + epilogue_args, + hw_info, + scheduler}; + + size_t workspace_size = Gemm::get_workspace_size(args); + auto workspace = + torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, + std::nullopt, a.device()); + const cudaStream_t stream = get_current_cuda_stream(a.get_device_index()); + + auto can_implement_status = gemm_op.can_implement(args); + STD_TORCH_CHECK( + can_implement_status == cutlass::Status::kSuccess, + "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status); + + // Run the GEMM + auto status = gemm_op.initialize(args, workspace.data_ptr()); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Failed to initialize MXFP4 GEMM: status=", (int)status, + " workspace_size=", workspace_size, + " num_experts=", num_experts, " M=", M, " N=", N, " K=", K); + + status = gemm_op.run(args, workspace.data_ptr(), stream); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Failed to run MXFP4 GEMM"); +} + +template +void run_mxfp4_blockwise_scaled_group_mm( + torch::stable::Tensor& output, const torch::stable::Tensor& a, + const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + int32_t version_num = get_sm_version_num(); +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 + if (version_num >= 100 && version_num < 120) { + run_mxfp4_blockwise_scaled_group_mm_sm100( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + return; + } +#endif + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, + "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ", + version_num, ". Required capability: 100"); +} + +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 +constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte; +// E8M0 scale factors stored as uint8 +constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte; +#endif + +#define CHECK_TYPE(x, st, m) \ + STD_TORCH_CHECK(x.scalar_type() == st, \ + ": Inconsistency of torch::stable::Tensor type:", m) +#define CHECK_TH_CUDA(x, m) \ + STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.") +#define CHECK_CONTIGUOUS(x, m) \ + STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.") +#define CHECK_INPUT(x, st, m) \ + CHECK_TH_CUDA(x, m); \ + CHECK_CONTIGUOUS(x, m); \ + CHECK_TYPE(x, st, m) + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets) { +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 + // Input validation + CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a"); + CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b"); + // MXFP4 uses E8M0 scale factors (stored as uint8) + CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale"); + CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales"); + + STD_TORCH_CHECK( + a_blockscale.dim() == 2, + "expected a_blockscale to be of shape [num_experts, rounded_m," + " k // group_size], observed rank: ", + a_blockscale.dim()) + STD_TORCH_CHECK(b_blockscales.dim() == 3, + "expected b_blockscale to be of shape: " + " [num_experts, n, k // group_size], observed rank: ", + b_blockscales.dim()) + STD_TORCH_CHECK(problem_sizes.dim() == 2, + "problem_sizes must be a 2D tensor"); + STD_TORCH_CHECK(problem_sizes.size(1) == 3, + "problem_sizes must have the shape (num_experts, 3)"); + STD_TORCH_CHECK( + problem_sizes.size(0) == expert_offsets.size(0), + "Number of experts in problem_sizes must match expert_offsets"); + STD_TORCH_CHECK( + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, + "problem_sizes must be int32."); + + int M = static_cast(a.size(0)); + int N = static_cast(b.size(1)); + int E = static_cast(b.size(0)); + int K = static_cast(2 * b.size(2)); + + if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + run_mxfp4_blockwise_scaled_group_mm( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + } else { + run_mxfp4_blockwise_scaled_group_mm( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + } +#else + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, + "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with " + "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+."); +#endif +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm)); +} diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu new file mode 100644 index 000000000000..78e4eda0c012 --- /dev/null +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -0,0 +1,432 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * MXFP4 activation quantization kernel for MoE experts. + * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales + * over 32-element groups. + * + * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for: + * - Block size 32 (2 threads per SF in PACK16 mode) + * - E8M0 (power-of-two) scale factors + * - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128) + */ + +// MXFP4 requires PACK16 mode (16 elements per thread) so that +// 2 threads cover 32-element blocks. This requires CUDA >= 12.9. +// Must be defined before any header that (transitively) includes +// nvfp4_utils.cuh. +#define NVFP4_ENABLE_ELTS16 1 + +#include +#include +#include +#include + +#include +#include +#include "libtorch_stable/torch_utils.h" +#include "libtorch_stable/dispatch_utils.h" +#include "cuda_vec_utils.cuh" +#include "cuda_utils.h" + +#include "nvfp4_utils.cuh" +static_assert(CVT_FP4_ELTS_PER_THREAD == 16, + "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)"); + +#include "launch_bounds_utils.h" + +namespace vllm { + +// MXFP4 block size constants +static constexpr int MXFP4_SF_VEC_SIZE = 32; + +// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF +// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF +static constexpr int MXFP4_NUM_THREADS_PER_SF = + MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + +// MXFP4 quantization kernel for experts. +// Uses 32-element blocks with E8M0 (UE8M0) scale factors. +// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses +// SiLU(gate)*up before quantization. +template +__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) + mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in, + fp4_packed_t* out, uint32_t* SFout, + uint32_t* input_offset_by_experts, + uint32_t* output_scale_offset_by_experts, + int n_experts, bool low_latency) { + using PackedVec = PackedVec; + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile + int32_t const numKTiles = (numCols + 127) / 128; + + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; + + for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + globalIdx += gridDim.x * blockDim.x) { + int rowIdx = globalIdx / colsPerRow; + int colIdx = globalIdx % colsPerRow; + + int rowIdx_in_expert = 0; + int expert_idx = 0; + + if constexpr (SMALL_NUM_EXPERTS) { + for (int i = 0; i < n_experts; i++) { + uint32_t current_offset = __ldca(&input_offset_by_experts[i]); + uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]); + if (rowIdx >= current_offset && rowIdx < next_offset) { + rowIdx_in_expert = rowIdx - current_offset; + expert_idx = i; + break; + } + } + } else { + uint32_t local_offsets[17]; + for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) { + *reinterpret_cast(local_offsets) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start])); + *reinterpret_cast(local_offsets + 4) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 4])); + *reinterpret_cast(local_offsets + 8) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 8])); + *reinterpret_cast(local_offsets + 12) = + __ldca(reinterpret_cast( + &input_offset_by_experts[chunk_start + 12])); + local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]); + +#pragma unroll + for (int i = 0; i < 16; i++) { + if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) { + rowIdx_in_expert = rowIdx - local_offsets[i]; + expert_idx = chunk_start + i; + break; + } + } + } + } + + // Load input and optionally apply fused SiLU+Mul + int64_t inOffset = rowIdx * inColsPerRow + colIdx; + PackedVec in_vec = reinterpret_cast(in)[inOffset]; + PackedVec quant_input; + if constexpr (FUSE_SILU_MUL) { + PackedVec in_vec_up = + reinterpret_cast(in)[inOffset + colsPerRow]; + quant_input = compute_silu_mul(in_vec, in_vec_up); + } else { + quant_input = in_vec; + } + + // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2 + int64_t outOffset = rowIdx * colsPerRow + colIdx; + auto& out_pos = out[outOffset]; + + uint32_t* SFout_in_expert = + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles; + + // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks + auto sf_out = + cvt_quant_to_fp4_get_sf_out_offset( + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert); + + // Block E8M0 scales only; no extra tensor-level scale in this path + constexpr float SFScaleVal = 1.0f; + // UE8M0_SF=true for MXFP4 E8M0 scale factors + out_pos = + cvt_warp_fp16_to_fp4( + quant_input, SFScaleVal, sf_out); + } +} + +// Large M_topk variant using shared memory for expert offsets +template +__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) + mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in, + fp4_packed_t* out, uint32_t* SFout, + uint32_t* input_offset_by_experts, + uint32_t* output_scale_offset_by_experts, + int n_experts) { + using PackedVec = PackedVec; + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, + "Vec size is not matched."); + + // MXFP4: numKTiles = ceil(numCols / 128) + int32_t const numKTiles = (numCols + 127) / 128; + + extern __shared__ uint32_t shared_input_offsets[]; + + if constexpr (SMALL_NUM_EXPERTS) { + for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) { + shared_input_offsets[i] = input_offset_by_experts[i]; + } + } else { + for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) { + *reinterpret_cast(&shared_input_offsets[i]) = + *reinterpret_cast(&input_offset_by_experts[i]); + } + if (threadIdx.x == 0) { + shared_input_offsets[n_experts] = input_offset_by_experts[n_experts]; + } + } + + __syncthreads(); + + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; + + for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + globalIdx += gridDim.x * blockDim.x) { + int rowIdx = globalIdx / colsPerRow; + int colIdx = globalIdx % colsPerRow; + + int rowIdx_in_expert = 0; + int expert_idx = 0; + + // Binary search through experts using shared memory + int left = 0, right = n_experts - 1; + while (left <= right) { + int mid = (left + right) / 2; + uint32_t mid_offset = shared_input_offsets[mid]; + uint32_t next_offset = shared_input_offsets[mid + 1]; + + if (rowIdx >= mid_offset && rowIdx < next_offset) { + rowIdx_in_expert = rowIdx - mid_offset; + expert_idx = mid; + break; + } else if (rowIdx < mid_offset) { + right = mid - 1; + } else { + left = mid + 1; + } + } + + int64_t inOffset = rowIdx * inColsPerRow + colIdx; + PackedVec in_vec = reinterpret_cast(in)[inOffset]; + PackedVec quant_input; + if constexpr (FUSE_SILU_MUL) { + PackedVec in_vec_up = + reinterpret_cast(in)[inOffset + colsPerRow]; + quant_input = compute_silu_mul(in_vec, in_vec_up); + } else { + quant_input = in_vec; + } + + int64_t outOffset = rowIdx * colsPerRow + colIdx; + auto& out_pos = out[outOffset]; + + // MXFP4 has no global scale - only block-level E8M0 scale factors + constexpr float SFScaleVal = 1.0f; + + uint32_t* SFout_in_expert = + SFout + output_scale_offset_by_experts[expert_idx] * numKTiles; + + auto sf_out = + cvt_quant_to_fp4_get_sf_out_offset( + rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert); + + out_pos = + cvt_warp_fp16_to_fp4( + quant_input, SFScaleVal, sf_out); + } +} + +template +void mxfp4_quant_impl(void* output, void* output_scale, void* input, + void* input_offset_by_experts, + void* output_scale_offset_by_experts, int m_topk, int k, + int n_experts, cudaStream_t stream) { + int multiProcessorCount = + get_device_attribute(cudaDevAttrMultiProcessorCount, -1); + + int const workSizePerRow = k / ELTS_PER_THREAD; + int const totalWorkSize = m_topk * workSizePerRow; + dim3 block(std::min(workSizePerRow, 512)); + int const numBlocksPerSM = + vllm_runtime_blocks_per_sm(static_cast(block.x)); + dim3 grid(std::min(static_cast((totalWorkSize + block.x - 1) / block.x), + multiProcessorCount * numBlocksPerSM)); + while (grid.x <= multiProcessorCount && block.x > 64) { + grid.x *= 2; + block.x = (block.x + 1) / 2; + } + + int const blockRepeat = + (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x); + if (blockRepeat > 1) { + size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t); + if (n_experts >= 4) { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts); + } else { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts); + } + } else { + if (n_experts >= 16) { + mxfp4_cvt_fp16_to_fp4 + <<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts, /* bool low_latency */ true); + } else { + mxfp4_cvt_fp16_to_fp4<<>>( + m_topk, k, reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(output_scale), + reinterpret_cast(input_offset_by_experts), + reinterpret_cast(output_scale_offset_by_experts), + n_experts, /* bool low_latency */ true); + } + } +} + +} // namespace vllm + +/*Quantization entry for mxfp4 experts quantization*/ +#define CHECK_TH_CUDA(x, m) \ + STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x, m) \ + STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous") +#define CHECK_INPUT(x, m) \ + CHECK_TH_CUDA(x, m); \ + CHECK_CONTIGUOUS(x, m); + +constexpr auto HALF = torch::headeronly::ScalarType::Half; +constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16; +constexpr auto INT = torch::headeronly::ScalarType::Int; +constexpr auto UINT8 = torch::headeronly::ScalarType::Byte; + +static constexpr int MXFP4_BLOCK_SIZE = 32; + +static void validate_mxfp4_experts_quant_inputs( + torch::stable::Tensor const& output, + torch::stable::Tensor const& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts, int64_t m_topk, int64_t k) { + CHECK_INPUT(output, "output"); + CHECK_INPUT(output_scale, "output_scale"); + CHECK_INPUT(input, "input"); + CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts"); + CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts"); + + STD_TORCH_CHECK(output.dim() == 2); + STD_TORCH_CHECK(output_scale.dim() == 2); + STD_TORCH_CHECK(input.dim() == 2); + STD_TORCH_CHECK(input_offset_by_experts.dim() == 1); + STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1); + + STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16); + STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT); + STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT); + // output is uint8 (two mxfp4 values packed into one uint8) + // output_scale is int32 (four E8M0 values packed into one int32) + STD_TORCH_CHECK(output.scalar_type() == UINT8); + STD_TORCH_CHECK(output_scale.scalar_type() == INT); + + STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32"); + STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1); + STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1); + STD_TORCH_CHECK(output.size(0) == m_topk); + STD_TORCH_CHECK(output.size(1) == k / 2); + int scales_k = k / MXFP4_BLOCK_SIZE; + // K-dimension scale columns padded to a multiple of 4 for swizzle layout + int padded_k = (scales_k + (4 - 1)) / 4 * 4; + // 4 = 4 E8M0 values packed into one int32 + STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k); +} + +void mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + auto m_topk = input.size(0); + auto k = input.size(1); + + validate_mxfp4_experts_quant_inputs( + output, output_scale, input, input_offset_by_experts, + output_scale_offset_by_experts, n_experts, m_topk, k); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + input.scalar_type(), "mxfp4_experts_quant_kernel", [&] { + using cuda_type = vllm::CUDATypeConverter::Type; + vllm::mxfp4_quant_impl( + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, + stream); + }); +} + +void silu_and_mul_mxfp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts, + int64_t n_experts) { + auto m_topk = input.size(0); + auto k_times_2 = input.size(1); + STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)"); + auto k = k_times_2 / 2; + + validate_mxfp4_experts_quant_inputs( + output, output_scale, input, input_offset_by_experts, + output_scale_offset_by_experts, n_experts, m_topk, k); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] { + using cuda_type = vllm::CUDATypeConverter::Type; + vllm::mxfp4_quant_impl( + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), + input_offset_by_experts.data_ptr(), + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, + stream); + }); +} + +// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied +// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to +// .cpp files and cannot gate the registration from there. +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); + m.impl("silu_and_mul_mxfp4_experts_quant", + TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); +} diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu index f90bd543ab9f..744ae4f73112 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu @@ -277,7 +277,9 @@ void quant_impl(void* output, void* output_scale, void* input, (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x); if (blockRepeat > 1) { size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t); - if (n_experts >= 4) { + // The shared-memory vectorized offset load only handles full 4-expert + // chunks. Use the scalar specialization for the remainder cases. + if (n_experts >= 4 && n_experts % 4 == 0) { cvt_fp16_to_fp4 <<>>( m_topk, k, reinterpret_cast(input), @@ -299,7 +301,9 @@ void quant_impl(void* output, void* output_scale, void* input, n_experts); } } else { - if (n_experts >= 16) { + // The low-latency vectorized expert lookup only handles full 16-expert + // chunks. Fall back to the scalar lookup path for the remainder cases. + if (n_experts >= 16 && n_experts % 16 == 0) { cvt_fp16_to_fp4 <<>>( m_topk, k, reinterpret_cast(input), diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index 69b6564be750..8b0356815bb7 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -240,8 +240,9 @@ template __global__ void per_token_group_quant_8bit_packed_kernel( const T* __restrict__ input, void* __restrict__ output_q, unsigned int* __restrict__ output_s_packed, const int group_size, - const int num_groups, const int groups_per_block, const int groups_per_row, - const int mn, const int tma_aligned_mn, const float eps, + const int num_groups_padded, const int groups_per_block, + const int padded_groups_per_row, const int groups_per_row, const int mn, + const int tma_aligned_mn, const int num_scale_elems, const float eps, const float min_8bit, const float max_8bit) { const int threads_per_group = 16; const int64_t local_group_id = threadIdx.x / threads_per_group; @@ -249,51 +250,62 @@ __global__ void per_token_group_quant_8bit_packed_kernel( const int64_t block_group_id = blockIdx.x * groups_per_block; const int64_t global_group_id = block_group_id + local_group_id; - if (global_group_id >= num_groups) { + if (global_group_id >= num_groups_padded) { return; } - const int64_t block_group_offset = global_group_id * group_size; + // map flat group id to 2D indices (mn_idx, sf_k_idx) + const int sf_k_idx = + static_cast(global_group_id % padded_groups_per_row); + const int mn_idx = static_cast(global_group_id / padded_groups_per_row); - const T* group_input = input + block_group_offset; - DST_DTYPE* group_output = - static_cast(output_q) + block_group_offset; + // whether it is a valid group (not padding) + const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row); // shared memory to cache each group's data to avoid double DRAM reads. extern __shared__ __align__(16) char smem_raw[]; T* smem = reinterpret_cast(smem_raw); T* smem_group = smem + local_group_id * group_size; - const float y_s = - ComputeGroupScale(group_input, smem_group, group_size, lane_id, - threads_per_group, eps, max_8bit); - // pack 4 scales into a uint32 - if (lane_id == 0) { - // map flat group id to 2D indices (mn_idx, sf_k_idx) - const int sf_k_idx = static_cast(global_group_id % groups_per_row); - const int mn_idx = static_cast(global_group_id / groups_per_row); + // compute scale for valid groups + float y_s = 0.f; + if (is_valid_group) { + const T* group_input = + input + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + y_s = ComputeGroupScale(group_input, smem_group, group_size, + lane_id, threads_per_group, eps, max_8bit); + } - if (mn_idx < mn) { - // each uint32 in output_s_packed stores 4 packed scales - const int sf_k_pack_idx = sf_k_idx / 4; - const int pos = sf_k_idx % 4; + // pack 4 scales into a uint32 exponent + if (lane_id == 0) { + // each uint32 in output_s_packed stores 4 packed scales + const int sf_k_pack_idx = sf_k_idx / 4; + const int pos = sf_k_idx % 4; + const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; + if (is_valid_group) { // reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit // exponent, and place it into the correct byte of the 32-bit word. const unsigned int bits = __float_as_uint(y_s); - const unsigned int exponent = (bits >> 23u) & 0xffu; - const unsigned int contrib = exponent << (pos * 8u); - - const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; - // atomically OR 8-bit exponent into the packed scales buffer - atomicOr(output_s_packed + out_idx, contrib); + const uint8_t exponent = static_cast((bits >> 23u) & 0xffu); + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = exponent; + } else if (out_idx < num_scale_elems) { + // write zero for padding groups if within bounds of output_s_packed + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = 0; } } __syncthreads(); - QuantizeGroup(smem_group, group_output, group_size, lane_id, - threads_per_group, y_s, min_8bit, max_8bit); + if (is_valid_group) { + DST_DTYPE* group_output = + static_cast(output_q) + + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + QuantizeGroup(smem_group, group_output, group_size, lane_id, + threads_per_group, y_s, min_8bit, max_8bit); + } } void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, @@ -310,7 +322,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, const int64_t mn = input.numel() / k; const int64_t groups_per_row = k / group_size; - const int64_t num_groups = mn * groups_per_row; STD_TORCH_CHECK(output_s_packed.dim() == 2, "output_s_packed must be 2D, got dim=", output_s_packed.dim(), @@ -330,36 +341,46 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk, "], but got [", output_s_packed.size(0), ", ", output_s_packed.size(1), "]."); + // Verify column-major TMA-aligned layout + STD_TORCH_CHECK(output_s_packed.stride(0) == 1 && + output_s_packed.stride(1) == tma_aligned_mn, + "output_s_packed must have strides [1, ", tma_aligned_mn, + "], but got [", output_s_packed.stride(0), ", ", + output_s_packed.stride(1), "]."); cudaStream_t stream = get_current_cuda_stream(); constexpr int THREADS_PER_GROUP = 16; - const int groups_per_block = GetGroupsPerBlock(num_groups); + // Expand the grid to cover MN and K padding so every byte in + // output_s_packed is written (padding bytes get zeroed by the kernel). + const int64_t padded_groups_per_row = k_num_packed_sfk * 4; + const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row; + // Number of elements in output_s_packed. + const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn; + + const int groups_per_block = GetGroupsPerBlock(num_groups_padded); auto dst_type = output_q.scalar_type(); - const int num_blocks = num_groups / groups_per_block; + const int num_blocks = num_groups_padded / groups_per_block; const int num_threads = groups_per_block * THREADS_PER_GROUP; - // zero-initialize packed scales, since we use atomicOr to accumulate - // exponents from different groups. - torch::stable::zero_(output_s_packed); - -#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \ - do { \ - dim3 grid(num_blocks); \ - dim3 block(num_threads); \ - size_t smem_bytes = \ - static_cast(groups_per_block) * group_size * sizeof(T); \ - per_token_group_quant_8bit_packed_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - reinterpret_cast(output_s_packed.data_ptr()), \ - static_cast(group_size), static_cast(num_groups), \ - groups_per_block, static_cast(groups_per_row), \ - static_cast(mn), static_cast(tma_aligned_mn), \ - static_cast(eps), static_cast(min_8bit), \ - static_cast(max_8bit)); \ +#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \ + do { \ + dim3 grid(num_blocks); \ + dim3 block(num_threads); \ + size_t smem_bytes = \ + static_cast(groups_per_block) * group_size * sizeof(T); \ + per_token_group_quant_8bit_packed_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + static_cast(group_size), static_cast(num_groups_padded), \ + groups_per_block, static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(tma_aligned_mn), \ + static_cast(num_scale_elems), static_cast(eps), \ + static_cast(min_8bit), static_cast(max_8bit)); \ } while (0) VLLM_STABLE_DISPATCH_FLOATING_TYPES( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index c31844948e5f..124512e81628 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + // Compute NVFP4 block quantized tensor. ops.def( "scaled_fp4_quant(Tensor input," @@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," "Tensor output_scale_offset_by_experts) -> ()"); + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + // Fused SiLU+Mul+NVFP4 quantization. ops.def( "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " @@ -233,9 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("silu_and_mul_scaled_fp4_experts_quant", TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - - // W4A8 ops: impl registrations are in the source files - // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu) + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. #endif } diff --git a/csrc/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu new file mode 100644 index 000000000000..6245b02d6e98 --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -0,0 +1,879 @@ + +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include "cuda_compat.h" +#include "cuda_utils.h" +#include "core/registration.h" +#include "minimax_reduce_rms_kernel.h" + +#include + +#define FINAL_MASK 0xffffffff +#define MINIMAX_REDUCE_RMS_WARP_SIZE 32 + +namespace vllm { +namespace tensorrt_llm { + +template +struct LamportComm { + __device__ __forceinline__ LamportComm(void** workspace, int rank) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2]; + clear_ptr = &reinterpret_cast(workspace[NRanks * 3 + 1])[0]; + flag_value = *flag_ptr; + auto comm_size = reinterpret_cast(workspace[NRanks * 3 + 1])[1]; + clear_size = *clear_ptr; + int data_offset = flag_value % 3; + int clear_offset = (flag_value + 2) % 3; + for (int r = 0; r < NRanks; ++r) { + data_bufs[r] = reinterpret_cast(workspace[2 * NRanks + r]) + + data_offset * comm_size; + } + clear_buf = reinterpret_cast(workspace[2 * NRanks + rank]) + + clear_offset * comm_size; + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ void update(int64_t new_clear_size) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + while (*reinterpret_cast(counter_ptr) != gridDim.x) { + } + *flag_ptr = (flag_value + 1) % 3; + *clear_ptr = new_clear_size; + *counter_ptr = 0; + } + } + + int* counter_ptr; + int* flag_ptr; + int64_t* clear_ptr; + uint8_t* data_bufs[NRanks]; + uint8_t* clear_buf; + int64_t clear_size; + int flag_value; +}; + +__device__ __forceinline__ bool is_neg_zero(float v) { + return *reinterpret_cast(&v) == 0x80000000; +} + +__device__ __forceinline__ bool is_neg_zero(float4 v) { + return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) || + is_neg_zero(v.w); +} + +__device__ __forceinline__ float4 get_neg_zero() { + float4 vec; +#pragma unroll + for (int i = 0; i < 4; ++i) { + reinterpret_cast(&vec)[i] = 0x80000000; + } + return vec; +} + +template +__device__ __forceinline__ float rms_rsqrt(float& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(Dim); + v = rsqrtf((v * kInvDim) + eps); + return v; +} + +template +__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(Dim); + v.x = rsqrtf((v.x * kInvDim) + eps); + v.y = rsqrtf((v.y * kInvDim) + eps); + v.z = rsqrtf((v.z * kInvDim) + eps); + v.w = rsqrtf((v.w * kInvDim) + eps); + return v; +} +__device__ __forceinline__ float4 ld_global_volatile(float4* addr) { + float4 val; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w) + : "l"(addr)); + return val; +} + +__device__ __forceinline__ float ld_global_volatile(float* addr) { + float val; + asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(val) : "l"(addr)); + return val; +} + +// Used by the scalar (non-float4) kernel only +template +__inline__ __device__ T warpReduceSumV2(T* val) { +#pragma unroll + for (int i = 0; i < NUM; i++) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32); + } + return (T)(0.0f); +} + +template +__inline__ __device__ T blockReduceSumV2(T* val) { + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(val); + + if (lane == 0) { +#pragma unroll + for (int i = 0; i < NUM; i++) { + shared[i][wid] = val[i]; + } + } + + __syncthreads(); + + bool is_mask = threadIdx.x < (blockDim.x / 32.f); +#pragma unroll + for (int i = 0; i < NUM; i++) { + val[i] = is_mask ? shared[i][lane] : (T)(0.0f); + } + warpReduceSumV2(val); + return (T)0.0f; +} + +// for float4 version +template +__device__ __forceinline__ void local_warp_reduce_sum_array( + T* value_ptr, uint32_t active_mask = 0xffffffffu) { + static_assert(kNumThreads >= 1 && + kNumThreads <= MINIMAX_REDUCE_RMS_WARP_SIZE); +#pragma unroll + for (int i = 0; i < ArraySize; ++i) { +#pragma unroll + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { + value_ptr[i] += __shfl_xor_sync(active_mask, value_ptr[i], mask, + MINIMAX_REDUCE_RMS_WARP_SIZE); + } + } +} + +constexpr int next_pow2(int val) { + int result = 1; + while (result < val) { + result <<= 1; + } + return result; +} + +// --------------------------------------------------------------------------- + +template +class IndexHelper { + public: + __device__ __forceinline__ IndexHelper(MiniMaxReduceRMSParams const& params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + token_id = grid.cluster_rank(); + access_id_in_token = cluster.thread_rank(); + token_stride = grid.num_clusters(); +#else + token_id = blockIdx.x; + access_id_in_token = threadIdx.x; + token_stride = gridDim.x; +#endif + access_id = token_id * params.hidden_dim / kElemsPerAccess + + access_id_in_token; + access_stride = token_stride * params.hidden_dim / kElemsPerAccess; + tot_access = params.size_q / kElemsPerAccess; + } + + int token_id; + int access_id_in_token; + int token_stride; + int access_id; + int access_stride; + int tot_access; +}; + +/** +* this kernel is used to for minimax attention module +* input tensor [total_tokens, hidden_dim / tp_size], fp32 +* rms weight [hidden_dim / tp_size], bf16 +step 1: reduce from single rank to get the variance sum (reduce(input^2, +dim=-1)) step 2: reduce from all ranks to get the variance sum +(all_reduce(variance_sum)) step 3: calculate the rms norm (input * +rsqrt(variance + eps)) in this case, max hidden_dim is 6144 (float data), for +each token, we only need 6144 / 4 / tp_size = (1536 / tp_size) threads so we can +assume cluster size is 1 (tp_size >= 2) + */ +template +__global__ void __launch_bounds__(1024) + minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) { + IndexHelper index_helper(params); + int token_id = index_helper.token_id; + int access_id_in_token = index_helper.access_id_in_token; + int token_stride = index_helper.token_stride; + int access_id = index_helper.access_id; + int access_stride = index_helper.access_stride; + int tot_access = index_helper.tot_access; + int tot_tokens = params.size_q / params.hidden_dim; + float4 clear_vec = get_neg_zero(); + + LamportComm comm(params.workspace, params.rank); + int clear_access = comm.clear_size / kElemsPerAccess; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + for (int idx = access_id; idx < tot_access; + idx += access_stride, token_id += token_stride) { + alignas(16) DType vals[kElemsPerAccess]; + float sum_variance = 0.F; + *reinterpret_cast(vals) = + reinterpret_cast(params.allreduce_in)[idx]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + sum_variance += static_cast(vals[i]) * static_cast(vals[i]); + } + blockReduceSumV2(&sum_variance); + if (is_neg_zero(sum_variance)) { + sum_variance = 0.F; + } + if (threadIdx.x == 0) { + for (int r = 0; r < NRanks; ++r) { + reinterpret_cast( + comm.data_bufs[r])[(params.rank * tot_tokens) + token_id] = + (sum_variance); + } + } + + bool done = false; + float vars_all_ranks[NRanks]; + while (!done) { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + vars_all_ranks[r] = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(r * tot_tokens) + token_id]); + done &= !is_neg_zero(vars_all_ranks[r]); + } + } + sum_variance = 0.F; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + sum_variance += vars_all_ranks[r]; + } + + DType norm_weight[kElemsPerAccess]; + *reinterpret_cast::vec_type*>(norm_weight) = + reinterpret_cast::vec_type*>( + params.rms_gamma)[access_id_in_token]; + +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[i] = static_cast( + static_cast(vals[i]) * + rsqrtf( + (sum_variance / static_cast(params.hidden_dim) / NRanks) + + params.rms_eps) * + static_cast(norm_weight[i])); + } + + reinterpret_cast(params.rms_norm_out)[idx] = + *reinterpret_cast(vals); + } + for (int idx = access_id; idx < clear_access; idx += access_stride) { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + comm.update(params.size_q * NRanks); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +/** + * Float4 variant: process 4 rows at once, allreduce variance sums as float4 for + * better memory coalescing. sum_variance is always float; applies to all DTypes + * (half, bf16, float). When tot_tokens % 4 != 0, the last group pads rows with + * zeros; padded rows are not written to rms_norm_out. IsQK: when true, process + * Q+K in one loop with doubled comm buffer; when false, single-matrix (Q only). + */ +template +__global__ void __launch_bounds__(1024) + minimax_reduce_qk_rms_kernel_lamport_float4(MiniMaxReduceRMSParams params) { + // Compile-time per-rank dimensions + constexpr int RankQDim = OriginQDim / NRanks; + constexpr int RankKDim = OriginKDim / NRanks; + // Threads needed to cover one row of Q / K with float4 accesses + constexpr int ThreadsPerRowQ = RankQDim / kElemsPerAccess; + constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess; + // Number of warps dedicated to Q / K + constexpr int NumWarpQ = (ThreadsPerRowQ + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) / + MINIMAX_REDUCE_RMS_WARP_SIZE; + constexpr int NumWarpK = (ThreadsPerRowK + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) / + MINIMAX_REDUCE_RMS_WARP_SIZE; + + int tot_tokens = params.size_q / RankQDim; + int tot_groups = (tot_tokens + 3) / 4; // ceiling; last group may be partial + + // Memory strides for strided qkv tensors (elements -> float4-access units) + int access_stride_q = (params.stride_q > 0 ? params.stride_q : RankQDim) / + kElemsPerAccess; + int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) / + kElemsPerAccess; + // Output strides: default to contiguous (hidden_dim / hidden_dim_k) + int access_stride_q_out = + (params.stride_q_out > 0 ? params.stride_q_out : params.hidden_dim) / + kElemsPerAccess; + int access_stride_k_out = + (params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) / + kElemsPerAccess; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + int group_id = grid.cluster_rank(); + int access_id_in_token = cluster.thread_rank(); + int group_stride = grid.num_clusters(); +#else + int group_id = blockIdx.x; + int access_id_in_token = threadIdx.x; + int group_stride = gridDim.x; +#endif + + bool is_q = (access_id_in_token < NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE); + int k_thread_idx = + access_id_in_token - (NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE); + bool is_valid_q = (access_id_in_token < ThreadsPerRowQ); + bool is_valid_k = (k_thread_idx >= 0 && k_thread_idx < ThreadsPerRowK); + float4 clear_vec = get_neg_zero(); + + // Shared memory for two-level block reduction and scale broadcast + __shared__ float block_reduce_sum[4][MINIMAX_REDUCE_RMS_WARP_SIZE + 1]; + __shared__ float global_scale_q[4]; + __shared__ float global_scale_k[4]; + + LamportComm comm(params.workspace, params.rank); + + DType norm_weight[kElemsPerAccess]{}; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + if (is_q) { + if (is_valid_q) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::vec_type const*>( + params.rms_gamma)[access_id_in_token]; + } + } else { + if (is_valid_k) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::vec_type const*>( + params.rms_gamma_k)[k_thread_idx]; + } + } + + // Main loop: process one group of 4 tokens per iteration. + for (int g = group_id; g < tot_groups; g += group_stride) { + alignas(16) DType vals[4][kElemsPerAccess]{}; + float warp_sum_variance[4]{0.F, 0.F, 0.F, 0.F}; + + if (is_q) { +#pragma unroll + for (int row = 0; row < 4; ++row) { + int token_r = g * 4 + row; + if (token_r >= tot_tokens || !is_valid_q) { + continue; + } + int idx_r = token_r * access_stride_q + access_id_in_token; + *reinterpret_cast(&vals[row][0]) = + reinterpret_cast(params.allreduce_in)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(vals[row][i]); + warp_sum_variance[row] += x * x; + } + } + } else { +#pragma unroll + for (int row = 0; row < 4; ++row) { + int token_r = g * 4 + row; + if (token_r >= tot_tokens || !is_valid_k) { + continue; + } + int idx_r = token_r * access_stride_k + k_thread_idx; + *reinterpret_cast(&vals[row][0]) = + reinterpret_cast(params.allreduce_in_k)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(vals[row][i]); + warp_sum_variance[row] += x * x; + } + } + } + + local_warp_reduce_sum_array( + warp_sum_variance); + // Warp lane 0 writes its warp's partial sum to shared memory + int lane = threadIdx.x & (MINIMAX_REDUCE_RMS_WARP_SIZE - 1); + if (lane == 0) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + block_reduce_sum[t][threadIdx.x / MINIMAX_REDUCE_RMS_WARP_SIZE] = + warp_sum_variance[t]; + } + } + __syncthreads(); + + int tid = threadIdx.x; + + if (tid < MINIMAX_REDUCE_RMS_WARP_SIZE) { + constexpr int kNumWarpQPow2 = + (next_pow2(NumWarpQ) > NRanks) ? next_pow2(NumWarpQ) : NRanks; + float local_sum[4]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + local_sum[t] = (tid < NumWarpQ) ? block_reduce_sum[t][tid] : 0.F; + } + // After this, all kNumWarpQPow2 lanes (including tid 0..NRanks-1) have + // the total Q sum-of-squares for all 4 tokens. + local_warp_reduce_sum_array(local_sum); + + if (tid < NRanks) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + if (is_neg_zero(local_sum[t])) { + local_sum[t] = 0.F; + } + } + // Parallel push: thread tid writes this rank's Q sum to rank tid's buf + reinterpret_cast( + comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] = + *reinterpret_cast(local_sum); + + // Parallel pull: thread tid reads rank tid's contribution from + // this rank's (params.rank's) buffer + bool done = false; + float4 var_all_ranks; + while (!done) { + done = true; + var_all_ranks = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(tid * tot_groups * 2) + (2 * g)]); + done &= !is_neg_zero(var_all_ranks); + } + + // Warp-level allreduce: each of the NRanks threads holds one rank's + // partial sum; after this all NRanks threads have the global total. + constexpr uint32_t kQActiveMask = (1u << NRanks) - 1u; + local_warp_reduce_sum_array( + reinterpret_cast(&var_all_ranks), kQActiveMask); + + // Thread 0 computes rsqrt with compile-time Dim and writes to smem + if (tid == 0) { + *reinterpret_cast(global_scale_q) = + rms_rsqrt(var_all_ranks, params.rms_eps); + } + } + } else if (tid >= MINIMAX_REDUCE_RMS_WARP_SIZE * NumWarpQ && + tid < MINIMAX_REDUCE_RMS_WARP_SIZE * (NumWarpQ + 1)) { + // --- K leader warp --- + constexpr int kNumWarpKPow2 = + (next_pow2(NumWarpK) > NRanks) ? next_pow2(NumWarpK) : NRanks; + float local_sum[4]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + local_sum[t] = (k_thread_idx < NumWarpK) + ? block_reduce_sum[t][NumWarpQ + k_thread_idx] + : 0.F; + } + local_warp_reduce_sum_array(local_sum); + + if (k_thread_idx < NRanks) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + if (is_neg_zero(local_sum[t])) { + local_sum[t] = 0.F; + } + } + reinterpret_cast( + comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) + + (2 * g + 1)] = + *reinterpret_cast(local_sum); + + bool done = false; + float4 var_all_ranks; + while (!done) { + done = true; + var_all_ranks = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(k_thread_idx * tot_groups * 2) + + (2 * g + 1)]); + done &= !is_neg_zero(var_all_ranks); + } + + constexpr uint32_t kKActiveMask = (1u << NRanks) - 1u; + local_warp_reduce_sum_array( + reinterpret_cast(&var_all_ranks), kKActiveMask); + + if (k_thread_idx == 0) { + *reinterpret_cast(global_scale_k) = + rms_rsqrt(var_all_ranks, params.rms_eps); + } + } + } + __syncthreads(); + + if (is_q) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + warp_sum_variance[t] = global_scale_q[t]; + } +#pragma unroll + for (int r = 0; r < 4; ++r) { +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(norm_weight[i])); + } + int token_r = g * 4 + r; + if (token_r >= tot_tokens || !is_valid_q) { + continue; + } + int idx_out = token_r * access_stride_q_out + access_id_in_token; + reinterpret_cast(params.rms_norm_out)[idx_out] = + *reinterpret_cast(&vals[r][0]); + } + } else { +#pragma unroll + for (int t = 0; t < 4; ++t) { + warp_sum_variance[t] = global_scale_k[t]; + } +#pragma unroll + for (int r = 0; r < 4; ++r) { +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(norm_weight[i])); + } + int token_r = g * 4 + r; + if (token_r >= tot_tokens || !is_valid_k) { + continue; + } + int idx_out = token_r * access_stride_k_out + k_thread_idx; + reinterpret_cast(params.rms_norm_out_k)[idx_out] = + *reinterpret_cast(&vals[r][0]); + } + } + } // end group loop +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + int clear_access = static_cast(comm.clear_size / kElemsPerAccess); + int clear_stride = group_stride * blockDim.x; + for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access; + idx += clear_stride) { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + + comm.update(static_cast(2) * tot_groups * kElemsPerAccess * + NRanks); +} + +int get_sm_count() { + static int sm_count = 0; + if (sm_count == 0) { + int device_id; + CUDA_CHECK(cudaGetDevice(&device_id)); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + sm_count = device_prop.multiProcessorCount; + } + return sm_count; +} + +inline int getSMVersion(bool queryRealSmArch = false) { + int device{-1}; + CUDA_CHECK(cudaGetDevice(&device)); + int sm_major = 0; + int sm_minor = 0; + CUDA_CHECK(cudaDeviceGetAttribute(&sm_major, + cudaDevAttrComputeCapabilityMajor, device)); + CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor, + cudaDevAttrComputeCapabilityMinor, device)); + int sm = sm_major * 10 + sm_minor; + if (sm == 121 && !queryRealSmArch) { + return 120; + } + return sm; +} + +template +int get_max_active_blocks(KernelFunc kernel, int block_size, + int dynamic_smem = 0) { + int max_active = 0; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active, kernel, block_size, dynamic_smem)); + return std::max(max_active, 1); +} + +template +void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) { + static int SM = getSMVersion(); + int token_num = params.size_q / params.hidden_dim; + int sm_count = get_sm_count(); + int cluster_size = 1; + int cluster_num = token_num; + int threads_per_token = params.hidden_dim / kElemsPerAccess; + int block_size = threads_per_token; + + int max_blocks_per_sm = get_max_active_blocks( + minimax_reduce_rms_kernel_lamport, block_size); + int max_grid = max_blocks_per_sm * sm_count; + + int grid_size = + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) * + cluster_size; + + cudaLaunchConfig_t cfg; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + + cudaLaunchAttribute attribute[2]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = SM >= 90 ? 2 : 0; + + CUDA_CHECK(cudaLaunchKernelEx( + &cfg, minimax_reduce_rms_kernel_lamport, params)); +} + +template +void minimax_reduce_rms_kernel_launcher_float4( + MiniMaxReduceRMSParams const& params) { + TORCH_CHECK(params.size_q % params.hidden_dim == 0); + TORCH_CHECK(params.hidden_dim % kElemsPerAccess == 0); + if (params.stride_q > 0) { + TORCH_CHECK(params.stride_q % kElemsPerAccess == 0); + } + TORCH_CHECK(params.allreduce_in_k != nullptr, + "float4 QK kernel requires K input"); + TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k); + TORCH_CHECK(params.size_k % params.hidden_dim_k == 0); + TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess == 0); + TORCH_CHECK(params.size_q / params.hidden_dim == + params.size_k / params.hidden_dim_k); + if (params.stride_k > 0) { + TORCH_CHECK(params.stride_k % kElemsPerAccess == 0); + } + + int token_num = params.size_q / params.hidden_dim; + int tot_groups = (token_num + 3) / 4; + if (tot_groups == 0) { + return; + } + + static int SM = getSMVersion(); + int sm_count = get_sm_count(); + int cluster_size = 1; + int cluster_num = tot_groups; + + int access_per_row_q = params.hidden_dim / kElemsPerAccess; + int access_per_row_k = params.hidden_dim_k / kElemsPerAccess; + + // Round each section up to a warp boundary + auto divUp = [](int a, int b) { return (a + b - 1) / b * b; }; + int block_size = divUp(access_per_row_q, MINIMAX_REDUCE_RMS_WARP_SIZE) + + divUp(access_per_row_k, MINIMAX_REDUCE_RMS_WARP_SIZE); + + auto kfn = + minimax_reduce_qk_rms_kernel_lamport_float4; + + int max_blocks_per_sm = get_max_active_blocks(kfn, block_size); + int max_grid = max_blocks_per_sm * sm_count; + int grid_size = + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) * + cluster_size; + + cudaLaunchConfig_t cfg; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + + cudaLaunchAttribute attribute[2]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = SM >= 90 ? 2 : 0; + + CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params)); +} + +template +void dispatch_dtype(MiniMaxReduceRMSParams const& params) { + // Use the optimized QK float4 kernel when: + // - K input is present, AND + // - the full (NRanks * per-rank) dimensions match the MiniMax M2 shape. + // Otherwise fall back to the scalar kernel. + bool use_float4 = (params.allreduce_in_k != nullptr) && + (params.hidden_dim * params.nranks == 6144) && + (params.hidden_dim_k * params.nranks == 1024); + + if (params.dtype == at::ScalarType::Half) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4( + params); + } else { + minimax_reduce_rms_kernel_launcher(params); + } + } else if (params.dtype == at::ScalarType::BFloat16) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144, + 1024>(params); + } else { + minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params); + } + } else if (params.dtype == at::ScalarType::Float) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4( + params); + } else { + minimax_reduce_rms_kernel_launcher(params); + } + } else { + TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op"); + } +} + +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { + if (params.nranks == 2) { + dispatch_dtype<2>(params); + } else if (params.nranks == 4) { + dispatch_dtype<4>(params); + } else if (params.nranks == 8) { + dispatch_dtype<8>(params); + } else if (params.nranks == 16) { + dispatch_dtype<16>(params); + } else { + TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!"); + } +} +} // namespace tensorrt_llm +} // namespace vllm + +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, + torch::Tensor const& norm_weight, + torch::Tensor workspace, int64_t const rank, + int64_t const nranks, double const eps) { + auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); + + allreduce_params.nranks = static_cast(nranks); + allreduce_params.rank = static_cast(rank); + allreduce_params.dtype = input.scalar_type(); + allreduce_params.size_q = static_cast(input.numel()); + allreduce_params.hidden_dim = static_cast(input.size(-1)); + allreduce_params.stride_q = allreduce_params.hidden_dim; + allreduce_params.workspace = + reinterpret_cast(workspace.mutable_data_ptr()); + allreduce_params.allreduce_in = input.data_ptr(); + allreduce_params.rms_gamma = norm_weight.data_ptr(); + allreduce_params.rms_eps = static_cast(eps); + allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + torch::Tensor rms_norm_out = torch::empty_like(input); + allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr(); + + vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params); + + return rms_norm_out; +} + +std::tuple minimax_allreduce_rms_qk( + torch::Tensor qkv, torch::Tensor const& norm_weight_q, + torch::Tensor const& norm_weight_k, torch::Tensor workspace, + int64_t const q_size, int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps) { + TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D"); + TORCH_CHECK(qkv.is_contiguous(), + "minimax_allreduce_rms_qk: qkv must be contiguous"); + int64_t qkv_dim = qkv.size(-1); + TORCH_CHECK(qkv_dim == q_size + 2 * kv_size, + "minimax_allreduce_rms_qk: qkv last dim must equal " + "q_size + 2 * kv_size"); + TORCH_CHECK(rank < nranks, + "minimax_allreduce_rms_qk: rank must be less than nranks"); + + int64_t num_tokens = qkv.size(0); + int elem_bytes = qkv.element_size(); + + torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options()); + torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options()); + + auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); + params.nranks = static_cast(nranks); + params.rank = static_cast(rank); + params.dtype = qkv.scalar_type(); + params.size_q = static_cast(num_tokens * q_size); + params.hidden_dim = static_cast(q_size); + params.size_k = static_cast(num_tokens * kv_size); + params.hidden_dim_k = static_cast(kv_size); + params.stride_q = static_cast(qkv_dim); + params.stride_k = static_cast(qkv_dim); + params.stride_q_out = 0; // q_out is contiguous; kernel uses hidden_dim + params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k + params.workspace = reinterpret_cast(workspace.mutable_data_ptr()); + + uint8_t* base = static_cast(qkv.data_ptr()); + params.allreduce_in = base; + params.allreduce_in_k = base + q_size * elem_bytes; + params.rms_gamma = norm_weight_q.data_ptr(); + params.rms_gamma_k = norm_weight_k.data_ptr(); + params.rms_eps = static_cast(eps); + params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + + params.rms_norm_out = q_out.mutable_data_ptr(); + params.rms_norm_out_k = k_out.mutable_data_ptr(); + + vllm::tensorrt_llm::minimax_reduce_rms_op(params); + return {q_out, k_out}; +} diff --git a/csrc/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h new file mode 100644 index 000000000000..e8c2d012247b --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.h @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include + +namespace vllm { +namespace tensorrt_llm { + +template +struct ElemsPerAccess; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 4; + using vec_type = float4; +}; + +template +static constexpr int kElemsPerAccess = ElemsPerAccess::value; + +struct MiniMaxReduceRMSParams { + int nranks{}; + int rank{}; + at::ScalarType dtype{at::ScalarType::Undefined}; + int size_q{}; + int hidden_dim{}; + int size_k{}; + int hidden_dim_k{}; + int stride_q{}; // row stride for q input (elements); when > hidden_dim, + // q is part of a wider qkv tensor + int stride_k{}; // row stride for k input (elements); when > hidden_dim_k, + // k is part of a wider qkv tensor + int stride_q_out{}; // row stride for q output (elements); 0 = contiguous + int stride_k_out{}; // row stride for k output (elements); 0 = contiguous + void** workspace{}; + void* allreduce_in{}; + void* rms_norm_out{}; + void* rms_gamma{}; + void* allreduce_in_k{}; + void* rms_norm_out_k{}; + void* rms_gamma_k{}; + float rms_eps{}; + cudaStream_t stream{}; +}; + +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params); + +} // namespace tensorrt_llm +} // namespace vllm diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index d8d962887dab..973190935dfb 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -12,6 +12,15 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices, torch::Tensor& gating_output, bool renormalize, std::optional bias); +void topk_softplus_sqrt(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output, bool renormalize, + double routed_scaling_factor, + const c10::optional& correction_bias, + const c10::optional& input_ids, + const c10::optional& tid2eid); + void moe_sum(torch::Tensor& input, torch::Tensor& output); void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu index 833036da528e..57461a044f9d 100644 --- a/csrc/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -126,7 +126,9 @@ __launch_bounds__(TPB) __global__ { const int idx = thread_row_offset + ii; const float val = toFloat(input[idx]); - const float softmax_val = expf(val - float_max) * normalizing_factor; + float softmax_val = expf(val - float_max) * normalizing_factor; + // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream. + if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f; output[idx] = softmax_val; } } @@ -147,7 +149,9 @@ __launch_bounds__(TPB) __global__ { const int idx = thread_row_offset + ii; const float val = toFloat(input[idx]); - const float sigmoid_val = 1.0f / (1.0f + __expf(-val)); + float sigmoid_val = 1.0f / (1.0f + __expf(-val)); + // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream. + if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f; output[idx] = sigmoid_val; } } @@ -442,6 +446,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ } } + // Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs. + // NaN gating (from degenerate hidden states in CUDA graph padding) causes + // softmax to produce all-NaN, which makes the argmax loop always pick + // expert 0 for every top-k slot, producing duplicate expert IDs that + // crash FlashInfer's three-step MoE sort. + // With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1]. +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) { + row_chunk[ii] = 0.f; + } + } + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; // If bias is not null, use biased value for selection diff --git a/csrc/moe/topk_softplus_sqrt_kernels.cu b/csrc/moe/topk_softplus_sqrt_kernels.cu new file mode 100644 index 000000000000..50a8540a7374 --- /dev/null +++ b/csrc/moe/topk_softplus_sqrt_kernels.cu @@ -0,0 +1,715 @@ +/* + * Adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu + * Copyright (c) 2024, The vLLM team. + * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include +#include "../cuda_compat.h" +#include "../cub_helpers.h" +#ifndef USE_ROCM + #include + #include +#else + #include + #include +typedef __hip_bfloat16 __nv_bfloat16; +typedef __hip_bfloat162 __nv_bfloat162; +#endif + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +namespace vllm { +namespace moe { + +/// Aligned array type +template +struct alignas(Alignment) AlignedArray { + T data[N]; +}; + +template +__device__ __forceinline__ float toFloat(T value) { + if constexpr (std::is_same_v) { + return value; + } else if constexpr (std::is_same_v) { + return __bfloat162float(value); + } else if constexpr (std::is_same_v) { + return __half2float(value); + } +} + +#define FINAL_MASK 0xffffffff +template +__inline__ __device__ T warpReduceSum(T val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val += __shfl_xor_sync(FINAL_MASK, val, mask, 32); + return val; +} + +// ====================== TopK softplus_sqrt things +// =============================== + +/* + A Top-K gating softplus_sqrt written to exploit when the number of experts in + the MoE layers are a small power of 2. This allows us to cleanly share the + rows among the threads in a single warp and eliminate communication between + warps (so no need to use shared mem). + + It fuses the sigmoid, max and argmax into a single kernel. + + Limitations: + 1) This implementation is optimized for when the number of experts is a small + power of 2. Additionally it also supports when number of experts is multiple + of 64 which is still faster than the computing sigmoid and topK separately + (only tested on CUDA yet). 2) This implementation assumes k is small, but will + work for any k. +*/ + +template +__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ + void topkGatingSoftplusSqrt( + const InputType* input, const bool* finished, float* output, + const int num_rows, IndType* indices, int* source_rows, const int k, + const int start_expert, const int end_expert, const bool renormalize, + double routed_scaling_factor, const float* correction_bias, + const IndType* input_ids, const IndType* tid2eid) { + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v, + "InputType must be float, __nv_bfloat16, or __half"); + + // We begin by enforcing compile time assertions and setting up compile time + // constants. + static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), + "BYTES_PER_LDG must be power of 2"); + static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); + + // Number of bytes each thread pulls in per load + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType); + static constexpr int ELTS_PER_ROW = NUM_EXPERTS; + static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; + static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; + + if constexpr (std::is_same_v || + std::is_same_v) { + static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0, + "ELTS_PER_LDG must be 1 or even for 16-bit conversion"); + } + + // Restrictions based on previous section. + static_assert( + VPT % ELTS_PER_LDG == 0, + "The elements per thread must be a multiple of the elements per ldg"); + static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, + "The threads per row must cleanly divide the threads per warp"); + static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), + "THREADS_PER_ROW must be power of 2"); + static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, + "THREADS_PER_ROW can be at most warp size"); + + // We have NUM_EXPERTS elements per row. We specialize for small #experts + static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT; + static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; + static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; + + // Restrictions for previous section. + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, + "The elts per row must cleanly divide the total elt per warp"); + + // ===================== From this point, we finally start computing run-time + // variables. ======================== + + // Compute CTA and warp rows. We pack multiple rows into a single warp, and a + // block contains WARPS_PER_CTA warps. This, each block processes a chunk of + // rows. We start by computing the start row for each block. + const int cta_base_row = blockIdx.x * ROWS_PER_CTA; + + // Now, using the base row per thread block, we compute the base row per warp. + const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; + + // The threads in a warp are split into sub-groups that will work on a row. + // We compute row offset for each thread sub-group + const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; + const int thread_row = warp_base_row + thread_row_in_warp; + + // Threads with indices out of bounds should early exit here. + if (thread_row >= num_rows) { + return; + } + const bool row_is_active = finished ? !finished[thread_row] : true; + + // We finally start setting up the read pointers for each thread. First, each + // thread jumps to the start of the row it will read. + const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW; + + // Now, we compute the group each thread belong to in order to determine the + // first column to start loads. + const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; + const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; + const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; + + // Finally, we pull in the data from global mem + float row_chunk[VPT]; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert + // to float + if constexpr (std::is_same_v) { + using VecType = AlignedArray; + VecType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk); + const VecType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + } + } else if constexpr (std::is_same_v) { + if constexpr (ELTS_PER_LDG >= 2) { + using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>; + float2* row_chunk_f2 = reinterpret_cast(row_chunk); + const VecType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + int base_idx_f2 = ii * ELTS_PER_LDG / 2; +#pragma unroll + for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) { + row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2( + *reinterpret_cast(vec.data + jj * 2)); + } + } + } else { // ELTS_PER_LDG == 1 +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + const __nv_bfloat16* scalar_ptr = + thread_read_ptr + ii * THREADS_PER_ROW; + row_chunk[ii] = __bfloat162float(*scalar_ptr); + } + } + } else if constexpr (std::is_same_v) { + if constexpr (ELTS_PER_LDG >= 2) { + using VecType = AlignedArray<__half, ELTS_PER_LDG>; + float2* row_chunk_f2 = reinterpret_cast(row_chunk); + const VecType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + int base_idx_f2 = ii * ELTS_PER_LDG / 2; +#pragma unroll + for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) { + row_chunk_f2[base_idx_f2 + jj] = __half22float2( + *reinterpret_cast(vec.data + jj * 2)); + } + } + } else { // ELTS_PER_LDG == 1 +#pragma unroll + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW; + row_chunk[ii] = __half2float(*scalar_ptr); + } + } + } + constexpr float threshold = 20.0f; + constexpr float beta = 1.0f; + + // Hash MoE path: indices are predetermined from lookup table + if constexpr (USE_HASH) { + const IndType token_id = input_ids[thread_row]; + const IndType* expert_indices_for_token = tid2eid + token_id * k; +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + row_chunk[ii] = sqrtf(val); + } + float selected_sum = 0.f; +#pragma unroll + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int expert = expert_indices_for_token[k_idx]; + const int idx = k * thread_row + k_idx; + for (int ii = 0; ii < VPT; ++ii) { + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + if (expert == expert_idx) { + indices[idx] = expert; + selected_sum += row_chunk[ii]; + break; + } + } + } + // Compute per-thread scale (using warp reduction when renormalizing). + if (renormalize) { + selected_sum = warpReduceSum(selected_sum); + } + float scale = static_cast(routed_scaling_factor); + if (renormalize) { + const float denom = selected_sum > 0.f ? selected_sum : 1.f; + scale /= denom; + } + +#pragma unroll + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int expert = expert_indices_for_token[k_idx]; + const int idx = k * thread_row + k_idx; + for (int ii = 0; ii < VPT; ++ii) { + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + if (expert == expert_idx) { + output[idx] = row_chunk[ii] * scale; + break; + } + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + return; + } + +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + // Compute softplus: log(1 + exp(val)) with numerical stability + // When val > threshold, softplus(x) ≈ x to avoid exp overflow + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + val = sqrtf(val); + if (correction_bias) { + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + val = val + correction_bias[expert_idx]; + } + row_chunk[ii] = val; + } + + // Original TopK path: find top-k experts by score + // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find + // the topk elements in each row, along with the max index. + int start_col = first_elt_read_by_thread; + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; + + float selected_sum = 0.f; + for (int k_idx = 0; k_idx < k; ++k_idx) { + // First, each thread does the local argmax + float max_val = row_chunk[0]; + int expert = start_col; +#pragma unroll + for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; + ++ldg, col += COLS_PER_GROUP_LDG) { +#pragma unroll + for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { + float val = row_chunk[ldg * ELTS_PER_LDG + ii]; + + // No check on the experts here since columns with the smallest index + // are processed first and only updated if > (not >=) + if (val > max_val) { + max_val = val; + expert = col + ii; + } + } + } + +// Now, we perform the argmax reduce. We use the butterfly pattern so threads +// reach consensus about the max. This will be useful for K > 1 so that the +// threads can agree on "who" had the max value. That thread can then blank out +// their max with -inf and the warp can run more iterations... +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + float other_max = + VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW); + int other_expert = + VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW); + + // We want lower indices to "win" in every thread so we break ties this + // way + if (other_max > max_val || + (other_max == max_val && other_expert < expert)) { + max_val = other_max; + expert = other_expert; + } + } + + // Write the max for this k iteration to global memory. + if (thread_group_idx == 0) { + // Add a guard to ignore experts not included by this node + const bool node_uses_expert = + expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + // The lead thread from each sub-group will write out the final results to + // global memory. (This will be a single) thread per row of the + // input/output matrices. + const int idx = k * thread_row + k_idx; + if (correction_bias != nullptr) { + max_val -= correction_bias[expert]; + } + output[idx] = max_val; + indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + source_rows[idx] = k_idx * num_rows + thread_row; + if (renormalize) { + selected_sum += max_val; + } + } + + // Finally, we clear the value in the thread with the current max if there + // is another iteration to run. + if (k_idx + 1 < k) { + const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG; + const int thread_to_clear_in_group = + (expert / ELTS_PER_LDG) % THREADS_PER_ROW; + + // Only the thread in the group which produced the max will reset the + // "winning" value to -inf. + if (thread_group_idx == thread_to_clear_in_group) { + const int offset_for_expert = expert % ELTS_PER_LDG; + // Safe to set to any negative value since row_chunk values must be + // between 0 and 1. + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = + -10000.f; + } + } + } + + // Apply renormalization and routed scaling factor to final weights. + if (thread_group_idx == 0) { + float scale = static_cast(routed_scaling_factor); + if (renormalize) { + const float denom = selected_sum > 0.f ? selected_sum : 1.f; + scale /= denom; + } + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * thread_row + k_idx; + output[idx] = output[idx] * scale; + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +namespace detail { +// Constructs some constants needed to partition the work across threads at +// compile time. +template +struct TopkConstants { + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType); + static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || + EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, + ""); + static constexpr int VECs_PER_THREAD = + MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM)); + static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG; + static constexpr int THREADS_PER_ROW = EXPERTS / VPT; + static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW; +}; +} // namespace detail + +#define DISPATCH_HASH(use_hash, USE_HASH, ...) \ + if (use_hash) { \ + const bool USE_HASH = true; \ + static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \ + __VA_ARGS__ \ + } else { \ + const bool USE_HASH = false; \ + static_assert(USE_HASH == false, \ + "USE_HASH must be compile-time constant"); \ + __VA_ARGS__ \ + } + +template +void topkGatingSoftplusSqrtLauncherHelper( + const InputType* input, const bool* finished, float* output, + IndType* indices, int* source_row, const int num_rows, const int k, + const int start_expert, const int end_expert, const bool renormalize, + double routed_scaling_factor, const float* correction_bias, + const bool use_hash, const IndType* input_ids, const IndType* tid2eid, + cudaStream_t stream) { + static constexpr int BYTES_PER_LDG = + MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); + using Constants = + detail::TopkConstants; + static constexpr int VPT = Constants::VPT; + static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; + const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; + const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; + dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); + DISPATCH_HASH(use_hash, USE_HASH, { + auto* kernel = + &topkGatingSoftplusSqrt; +#ifndef USE_ROCM + cudaLaunchConfig_t config = {}; + config.gridDim = num_blocks; + config.blockDim = block_dim; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, + indices, source_row, k, start_expert, end_expert, + renormalize, routed_scaling_factor, correction_bias, + input_ids, tid2eid); +#else + kernel<<>>( + input, finished, output, num_rows, indices, source_row, k, start_expert, + end_expert, renormalize, routed_scaling_factor, correction_bias, + input_ids, tid2eid); +#endif + }) +} + +#ifndef USE_ROCM + #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ + static_assert(WARP_SIZE == 32, \ + "Unsupported warp size. Only 32 is supported for CUDA"); \ + topkGatingSoftplusSqrtLauncherHelper( \ + gating_output, nullptr, topk_weights, topk_indices, \ + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ + routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ + stream); +#else + #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ + if (WARP_SIZE == 64) { \ + topkGatingSoftplusSqrtLauncherHelper( \ + gating_output, nullptr, topk_weights, topk_indices, \ + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ + routed_scaling_factor, correction_bias, use_hash, input_ids, \ + tid2eid, stream); \ + } else if (WARP_SIZE == 32) { \ + topkGatingSoftplusSqrtLauncherHelper( \ + gating_output, nullptr, topk_weights, topk_indices, \ + token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ + routed_scaling_factor, correction_bias, use_hash, input_ids, \ + tid2eid, stream); \ + } else { \ + assert(false && \ + "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ + } +#endif + +template +void topkGatingSoftplusSqrtKernelLauncher( + const InputType* gating_output, float* topk_weights, IndType* topk_indices, + int* token_expert_indices, const int num_tokens, const int num_experts, + const int topk, const bool renormalize, double routed_scaling_factor, + const float* correction_bias, const bool use_hash, const IndType* input_ids, + const IndType* tid2eid, cudaStream_t stream) { + static constexpr int WARPS_PER_TB = 4; + static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; +#ifndef USE_ROCM + // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts + // elements can be loaded by a warp + static constexpr int BYTES_PER_LDG_MULTIPLE_64 = + (std::is_same_v || + std::is_same_v) + ? 4 + : 8; +#endif + switch (num_experts) { + case 1: + LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 2: + LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 4: + LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 8: + LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 16: + LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 32: + LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 64: + LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 128: + LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 256: + LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + case 512: + LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); + break; + // (CUDA only) support multiples of 64 when num_experts is not power of 2. + // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of + // num_experts, alternatively we can test 4 bytes loading and enable it in + // future. +#ifndef USE_ROCM + case 192: + LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 320: + LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 384: + LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 448: + LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; + case 576: + LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64); + break; +#endif + default: { + TORCH_CHECK(false, "Unsupported expert number: ", num_experts); + } + } +} + +} // namespace moe +} // namespace vllm + +template +void dispatch_topk_softplus_sqrt_launch( + const ComputeType* gating_output, torch::Tensor& topk_weights, + torch::Tensor& topk_indices, torch::Tensor& token_expert_indices, + int num_tokens, int num_experts, int topk, bool renormalize, + double routed_scaling_factor, + const c10::optional& correction_bias, + const c10::optional& input_ids, + const c10::optional& tid2eid, cudaStream_t stream) { + const float* bias_ptr = nullptr; + if (correction_bias.has_value()) { + bias_ptr = correction_bias.value().data_ptr(); + } + bool use_hash = false; + if (tid2eid.has_value()) { + TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE"); + use_hash = true; + } + if (topk_indices.scalar_type() == at::ScalarType::Int) { + const int* input_ids_ptr = nullptr; + const int* tid2eid_ptr = nullptr; + if (tid2eid.has_value()) { + input_ids_ptr = input_ids.value().data_ptr(); + tid2eid_ptr = tid2eid.value().data_ptr(); + } + + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.data_ptr(), + topk_indices.data_ptr(), token_expert_indices.data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream); + } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) { + const uint32_t* input_ids_ptr = nullptr; + const uint32_t* tid2eid_ptr = nullptr; + if (tid2eid.has_value()) { + input_ids_ptr = input_ids.value().data_ptr(); + tid2eid_ptr = tid2eid.value().data_ptr(); + } + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.data_ptr(), + topk_indices.data_ptr(), token_expert_indices.data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream); + } else { + TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long); + + const int64_t* input_ids_ptr = nullptr; + const int64_t* tid2eid_ptr = nullptr; + if (tid2eid.has_value()) { + input_ids_ptr = input_ids.value().data_ptr(); + tid2eid_ptr = tid2eid.value().data_ptr(); + } + + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.data_ptr(), + topk_indices.data_ptr(), token_expert_indices.data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream); + } +} + +void topk_softplus_sqrt( + torch::Tensor& topk_weights, // [num_tokens, topk] + torch::Tensor& topk_indices, // [num_tokens, topk] + torch::Tensor& token_expert_indices, // [num_tokens, topk] + torch::Tensor& gating_output, // [num_tokens, num_experts] + bool renormalize, double routed_scaling_factor, + const c10::optional& correction_bias, + const c10::optional& input_ids, + const c10::optional& tid2eid) { + const int num_experts = gating_output.size(-1); + const auto num_tokens = gating_output.numel() / num_experts; + const int topk = topk_weights.size(-1); + const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + if (gating_output.scalar_type() == at::ScalarType::Float) { + dispatch_topk_softplus_sqrt_launch( + gating_output.data_ptr(), topk_weights, topk_indices, + token_expert_indices, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, correction_bias, input_ids, tid2eid, stream); + } else if (gating_output.scalar_type() == at::ScalarType::Half) { + dispatch_topk_softplus_sqrt_launch<__half>( + reinterpret_cast(gating_output.data_ptr()), + topk_weights, topk_indices, token_expert_indices, num_tokens, + num_experts, topk, renormalize, routed_scaling_factor, correction_bias, + input_ids, tid2eid, stream); + } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) { + dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( + reinterpret_cast( + gating_output.data_ptr()), + topk_weights, topk_indices, token_expert_indices, num_tokens, + num_experts, topk, renormalize, routed_scaling_factor, correction_bias, + input_ids, tid2eid, stream); + } else { + TORCH_CHECK(false, "Unsupported gating_output data type: ", + gating_output.scalar_type()); + } +} \ No newline at end of file diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 7b627a6f8760..7bf56ba7a2cd 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -16,6 +16,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { "bias) -> ()"); m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid); +#ifndef USE_ROCM + m.def( + "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " + "token_expert_indices, Tensor gating_output, bool renormalize, float " + "routed_scaling_factor, Tensor? " + "bias, Tensor? input_ids, Tensor? tid2eid) -> ()"); + m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt); +#endif // Calculate the result of moe by summing up the partial results // from all selected experts. m.def("moe_sum(Tensor input, Tensor! output) -> ()"); diff --git a/csrc/nvfp4_kv_cache_kernels.cu b/csrc/nvfp4_kv_cache_kernels.cu new file mode 100644 index 000000000000..d6aa715c203d --- /dev/null +++ b/csrc/nvfp4_kv_cache_kernels.cu @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// NVFP4 KV cache store kernel. +// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them +// into the paged KV cache. +// +// Per page layout: [K_data | K_scale | V_data | V_scale] +// Both data and scale regions are contiguous per head, enabling direct +// TMA descriptor use. +// +// Reuses device functions from nvfp4_utils.cuh: +// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale +// - pack_fp4() for packing float pairs to fp4 +// - reciprocal_approximate_ftz() for fast reciprocal + +#define NVFP4_ENABLE_ELTS16 1 +#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" + +#include +#include +#include + +#include "dispatch_utils.h" + +namespace vllm { + +// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel. +// The swizzle pattern for HND layout is: +// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S] +// where T = block_size (page_size), S = scale_dim = head_size // 16. +// +// For a linear (t, s) position, the swizzled position is: +// swizzled_t = (t / 4) * 4 + (s / (S / 4)) +// swizzled_s = (s % (S / 4)) * 4 + (t % 4) +__device__ __forceinline__ int swizzle_scale_offset(int t, int s, + int scale_dim) { + int s_group = scale_dim / 4; + int swizzled_t = (t / 4) * 4 + (s / s_group); + int swizzled_s = (s % s_group) * 4 + (t % 4); + return swizzled_t * scale_dim + swizzled_s; +} + +// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache. +// +// Takes separate data and scale cache pointers for K and V. +// Within each KV side, data and scale are separate contiguous regions. +// +// Threading: one CUDA block per token, threads process heads and +// groups of 16 elements within each head. +template +__global__ void reshape_and_cache_nvfp4_kernel( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] + uint8_t* __restrict__ key_data_cache, // data region for K + uint8_t* __restrict__ value_data_cache, // data region for V + uint8_t* __restrict__ key_scale_cache, // scale region for K + uint8_t* __restrict__ value_scale_cache, // scale region for V + const int64_t* __restrict__ slot_mapping, // [num_actual_tokens] + const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale + const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale + const int64_t key_stride, // key.stride(0) in elements + const int64_t value_stride, // value.stride(0) in elements + const int num_heads, const int head_size, const int block_size, + const int64_t data_block_stride, // data cache stride for dim 0 + const int64_t data_head_stride, // data cache stride for heads + const int64_t data_block_offset_stride, // data cache stride for tokens + const int64_t scale_block_stride, // scale cache stride for dim 0 + const int64_t scale_head_stride, // scale cache stride for heads + const int64_t scale_block_offset_stride // scale cache stride for tokens +) { + using CudaType = typename CUDATypeConverter::Type; + using PVec = PackedVec; + + static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8 + static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS; + + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + if (slot_idx < 0) return; + + const int64_t block_idx = slot_idx / block_size; + const int block_offset = static_cast(slot_idx % block_size); + + const int scale_dim = head_size / 16; + const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; + + const int total_groups = num_heads * groups_per_head; + const int tid = threadIdx.x; + const int num_thread_groups = blockDim.x / THREADS_PER_SF; + const int tg_id = tid / THREADS_PER_SF; + const int tg_lane = tid % THREADS_PER_SF; + + // Process both K (kv=0) and V (kv=1) +#pragma unroll + for (int kv = 0; kv < 2; kv++) { + const scalar_t* __restrict__ src = (kv == 0) ? key : value; + const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr); + const int64_t src_stride = (kv == 0) ? key_stride : value_stride; + uint8_t* __restrict__ data_cache = + (kv == 0) ? key_data_cache : value_data_cache; + uint8_t* __restrict__ sc_cache = + (kv == 0) ? key_scale_cache : value_scale_cache; + + // Source pointer for this token (use actual stride, not assumed contiguous) + const CudaType* __restrict__ token_src = + reinterpret_cast(src) + token_idx * src_stride; + + // Destination bases in data and scale caches for this token's block + uint8_t* __restrict__ data_block = + data_cache + block_idx * data_block_stride; + uint8_t* __restrict__ scale_block = + sc_cache + block_idx * scale_block_stride; + + for (int g = tg_id; g < total_groups; g += num_thread_groups) { + const int head = g / groups_per_head; + const int group_in_head = g % groups_per_head; + + // Load 16 (or 8) bf16 elements from source + PVec in_vec; + const CudaType* __restrict__ src_ptr = + token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE + + tg_lane * ELTS; + +#pragma unroll + for (int i = 0; i < ELTS / 2; i++) { + in_vec.elts[i] = reinterpret_cast< + const typename PackedTypeConverter::Type*>(src_ptr)[i]; + } + + // Quantize: produces packed fp4 and writes scale factor. + uint8_t sf_val; + uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr; + + fp4_packed_t packed = cvt_warp_fp16_to_fp4( + in_vec, global_scale, sf_out_ptr); + + // Write packed FP4 data to data cache + uint8_t* __restrict__ data_dst = data_block + head * data_head_stride + + block_offset * data_block_offset_stride; + +#if CVT_FP4_PACK16 + { + // 16 elements → 8 bytes (u32x2) + int data_byte_offset = group_in_head * 8; + reinterpret_cast(data_dst + data_byte_offset)[0] = + (uint64_t(packed.hi) << 32) | uint64_t(packed.lo); + } +#else + { + // 8 elements → 4 bytes (uint32_t) + int data_byte_offset = + group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2; + reinterpret_cast(data_dst + data_byte_offset)[0] = packed; + } +#endif + + // Write block scale to scale cache. + // K (kv==0): linear layout (no swizzle). + // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel. + if (sf_out_ptr != nullptr) { + int scale_idx = group_in_head; + uint8_t* __restrict__ scale_dst; + if (kv == 0) { + scale_dst = scale_block + head * scale_head_stride + + block_offset * scale_block_offset_stride + scale_idx; + } else { + int swizzled_offset = + swizzle_scale_offset(block_offset, scale_idx, scale_dim); + int swizzled_t = swizzled_offset / scale_dim; + int swizzled_s = swizzled_offset % scale_dim; + scale_dst = scale_block + head * scale_head_stride + + swizzled_t * scale_block_offset_stride + swizzled_s; + } + *scale_dst = sf_val; + } + } + } +} + +} // namespace vllm + +// Non-template entry point callable from cache_kernels.cu. +// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1]. +// Each KV side contains both data and scale: +// page = [K_data | K_scale | V_data | V_scale] +void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + torch::Tensor& k_scale, + torch::Tensor& v_scale) { + int num_tokens = slot_mapping.size(0); + int num_heads = key.size(1); + int head_size = key.size(2); + int data_dim = head_size / 2; + int scale_dim = head_size / 16; + int full_dim = data_dim + scale_dim; + + // key_cache is kv_cache[:, 0] with shape + // [num_blocks, block_size, num_heads, full_dim] in logical order. + // Strides encode the physical layout (HND or NHD). + TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); + TORCH_CHECK(key_cache.size(3) == full_dim, + "key_cache last dim must be data_dim + scale_dim, got ", + key_cache.size(3), " expected ", full_dim); + + int block_size = key_cache.size(1); + + TORCH_CHECK(head_size % 16 == 0, + "head_size must be divisible by 16 for NVFP4 KV cache"); + TORCH_CHECK(block_size % 4 == 0, + "block_size must be divisible by 4 for NVFP4 KV cache swizzle"); + + // Detect physical layout from strides (based on full_dim). + // HND: head stride > block_offset stride. + bool is_hnd = key_cache.stride(2) > key_cache.stride(1); + + int64_t data_block_stride = key_cache.stride(0); // page_bytes + int64_t data_head_stride, data_block_offset_stride; + if (is_hnd) { + data_head_stride = (int64_t)block_size * data_dim; + data_block_offset_stride = data_dim; + } else { + data_head_stride = data_dim; + data_block_offset_stride = (int64_t)num_heads * data_dim; + } + + // Page layout: [K_data | K_scale | V_data | V_scale] + // Scale follows data within each KV side. + int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim; + + uint8_t* key_scale_ptr = key_cache.data_ptr() + data_per_kv; + uint8_t* value_scale_ptr = value_cache.data_ptr() + data_per_kv; + + // Scale strides: same page stride, inner strides from layout. + int64_t scale_block_stride = data_block_stride; + int64_t scale_head_stride, scale_block_offset_stride; + if (is_hnd) { + scale_head_stride = (int64_t)block_size * scale_dim; + scale_block_offset_stride = scale_dim; + } else { + scale_head_stride = scale_dim; + scale_block_offset_stride = (int64_t)num_heads * scale_dim; + } + + const float* k_scale_ptr = k_scale.data_ptr(); + const float* v_scale_ptr = v_scale.data_ptr(); + + int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; + int total_groups = num_heads * groups_per_head; + constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + int num_threads = std::min(total_groups * THREADS_PER_SF, 512); + num_threads = ((num_threads + 31) / 32) * 32; + + dim3 grid(num_tokens); + dim3 block(num_threads); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_REDUCED_FLOATING_TYPES( + key.scalar_type(), "reshape_and_cache_nvfp4", [&] { + vllm::reshape_and_cache_nvfp4_kernel + <<>>( + key.data_ptr(), value.data_ptr(), + key_cache.data_ptr(), value_cache.data_ptr(), + key_scale_ptr, value_scale_ptr, + slot_mapping.data_ptr(), k_scale_ptr, v_scale_ptr, + key.stride(0), value.stride(0), num_heads, head_size, + block_size, data_block_stride, data_head_stride, + data_block_offset_stride, scale_block_stride, scale_head_stride, + scale_block_offset_stride); + }); +} diff --git a/csrc/ops.h b/csrc/ops.h index cc58422231ff..16a78f570cf6 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -96,7 +97,13 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q, int64_t num_heads_k, int64_t num_heads_v, int64_t head_dim, double eps, torch::Tensor& q_weight, torch::Tensor& k_weight, torch::Tensor& cos_sin_cache, - bool is_neox, torch::Tensor& position_ids); + bool is_neox, torch::Tensor& position_ids, + int64_t forced_token_heads_per_warp); + +void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache, + torch::Tensor const& slot_mapping, torch::Tensor const& position_ids, + torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size); void apply_repetition_penalties_(torch::Tensor& logits, const torch::Tensor& prompt_mask, @@ -114,9 +121,9 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, int64_t numRows, int64_t stride0, int64_t stride1, int64_t topK); -void large_context_topk(const torch::Tensor& score, torch::Tensor& indices, - const torch::Tensor& lengths, - std::optional row_starts_opt); +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); void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, torch::Tensor& scale, @@ -143,20 +150,21 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input, std::optional residual, int64_t group_size, bool is_scale_transposed); -#ifndef USE_ROCM void silu_and_mul_per_block_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scales, int64_t group_size, std::optional scale_ub, bool is_scale_transposed); -#endif void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, std::optional key, int64_t head_size, - torch::Tensor& cos_sin_cache, bool is_neox); + torch::Tensor& cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse); 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); @@ -310,4 +318,16 @@ int64_t qr_max_size(); #ifndef USE_ROCM void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b); -#endif \ No newline at end of file +#endif + +#ifndef USE_ROCM +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, + torch::Tensor const& norm_weight, + torch::Tensor workspace, int64_t const rank, + int64_t const nranks, double const eps); +std::tuple minimax_allreduce_rms_qk( + torch::Tensor qkv, torch::Tensor const& norm_weight_q, + torch::Tensor const& norm_weight_k, torch::Tensor workspace, + int64_t const q_size, int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh new file mode 100644 index 000000000000..8b9d10ff83dd --- /dev/null +++ b/csrc/persistent_topk.cuh @@ -0,0 +1,1309 @@ +/* + * Persistent TopK Scheduler for DSA Indexer + */ + +#ifndef PERSISTENT_TOPK_CUH_ +#define PERSISTENT_TOPK_CUH_ + +#include +#include +#include +#include +#include + +namespace vllm { +namespace persistent { + +// ============================================================================ +// Constants +// ============================================================================ + +constexpr int kThreadsPerBlock = 1024; +constexpr int RADIX = 256; + +// Medium path: all shared state in dynamic smem (no static __shared__, +// which would inflate the kernel's smem footprint and kill occupancy +// for the decode/trivial paths). +constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 +constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 +constexpr size_t kMediumHeaderSize = + (kMediumHistBytes + kMediumScalarsBytes + 127) & ~size_t(127); // 3200 +constexpr int MAX_BUFFERED_ITEMS = 4096; +constexpr size_t kSmemMedium = + kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968 +constexpr uint32_t RADIX_THRESHOLD = 32768; + +// Decode path constants +constexpr int kDecodeBins = 2048; +constexpr uint32_t HIST2048_THRESHOLD = 8192; + +// Large path: fixed shared memory for histograms + scalars +constexpr size_t kFixedSmemLarge = + ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15); + +// ============================================================================ +// Common helpers +// ============================================================================ + +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return static_cast(key >> 8); +} + +// ============================================================================ +// Vectorized load helpers +// ============================================================================ + +// Unconditional float4 load with cache hint (.cg = cache at global level only). +__device__ __forceinline__ void load_float4(const float* ptr, float& v0, + float& v1, float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// Per-element predicated scalar loads with -inf default. +__device__ __forceinline__ void load_float4_predicated(const float* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + int p0 = (base < seq_len); + int p1 = (base + 1 < seq_len); + int p2 = (base + 2 < seq_len); + int p3 = (base + 3 < seq_len); + asm volatile( + "{\n" + " .reg .pred pr0, pr1, pr2, pr3;\n" + " setp.ne.u32 pr0, %4, 0;\n" + " setp.ne.u32 pr1, %5, 0;\n" + " setp.ne.u32 pr2, %6, 0;\n" + " setp.ne.u32 pr3, %7, 0;\n" + " mov.u32 %0, 0xFF800000;\n" + " mov.u32 %1, 0xFF800000;\n" + " mov.u32 %2, 0xFF800000;\n" + " mov.u32 %3, 0xFF800000;\n" + " @pr0 ld.global.cg.u32 %0, [%8];\n" + " @pr1 ld.global.cg.u32 %1, [%8+4];\n" + " @pr2 ld.global.cg.u32 %2, [%8+8];\n" + " @pr3 ld.global.cg.u32 %3, [%8+12];\n" + "}\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// ============================================================================ +// Large path: inter-CTA coordination state (one per group) +// ============================================================================ + +struct RadixRowState { + uint32_t histogram[3][256]; // Triple-buffered histograms + uint32_t remaining_k; + uint32_t prefix; + int arrival_counter; + int output_counter; +}; + +// ============================================================================ +// Kernel parameters +// ============================================================================ + +struct PersistentTopKParams { + const float* __restrict__ input; // [num_rows, stride] + int32_t* __restrict__ output; // [num_rows, top_k] + int32_t* __restrict__ lengths; // [num_rows] + RadixRowState* row_states; // large path: per-group state + uint32_t num_rows; + uint32_t stride; + uint32_t top_k; // actual k value for output stride + uint32_t chunk_size; // large path: elements per CTA + uint32_t ctas_per_group; // 1=medium, >1=large + uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit) +}; + +// ============================================================================ +// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) +// Uses 11-bit half-precision bins for fine granularity. +// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg. +// ============================================================================ + +// 11-bit bin from half-precision representation (ascending: high values -> high +// bins) +__device__ __forceinline__ uint32_t decode_bin(float x) { + __half hx = __float2half(x); + uint16_t bits = __half_as_ushort(hx); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> 5; +} + +template +__device__ __noinline__ void histogram_2048_topk( + const float* __restrict__ logits, int32_t* __restrict__ output_indices, + int32_t seq_len) { + extern __shared__ int decode_smem[]; + const int tx = threadIdx.x; + const int lane = tx & 31; + + // ---- Layout constants ---- + constexpr int SBASE = 8192 - 8; // 8184 + constexpr int RHIST = RADIX + 128; // 384 + constexpr int BOFF = 2 * RHIST; // 768 + constexpr int DBUF = (SBASE - BOFF) / 2; // 3708 + constexpr int MAX_ITEMS_PER_THREAD = + (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; + + enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; + + // ---- Initialize scalars (prevents stale data from prior rows) ---- + if (tx < 8) { + decode_smem[SBASE + tx] = 0; + } + + // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- + int* histo = decode_smem; + uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; + int nitems = 0; + + for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) { + histo[i] = 0; + } + __syncthreads(); + + const int n_vec = (seq_len + 3) >> 2; + const bool row_aligned = ((reinterpret_cast(logits) & 15) == 0); + + for (int i = tx; i < n_vec; i += kThreadsPerBlock) { + const int base = i << 2; + float v0, v1, v2, v3; + + if (row_aligned && base + 3 < seq_len) { + load_float4(logits + base, v0, v1, v2, v3); + } else { + load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3); + } + + const uint16_t b0 = static_cast(decode_bin(v0)); + const uint16_t b1 = static_cast(decode_bin(v1)); + const uint16_t b2 = static_cast(decode_bin(v2)); + const uint16_t b3 = static_cast(decode_bin(v3)); + reg_bins[nitems++] = b0; + reg_bins[nitems++] = b1; + reg_bins[nitems++] = b2; + reg_bins[nitems++] = b3; + atomicAdd(&histo[b0], 1); + atomicAdd(&histo[b1], 1); + atomicAdd(&histo[b2], 1); + atomicAdd(&histo[b3], 1); + } + __syncthreads(); + + // ---- CUB suffix sum ---- + using BlockScanT = cub::BlockScan; + const int h0 = histo[2 * tx]; + const int pair_sum = h0 + histo[2 * tx + 1]; + + auto& scan_storage = *reinterpret_cast( + decode_smem + kDecodeBins); + + int pair_prefix, total; + BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total); + + // Find threshold bin purely from registers + const int pair_suffix = total - pair_prefix; + + if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { + decode_smem[SBASE + sTHR] = 2 * tx; + } + { + const int right_suf = pair_suffix - h0; + const int next_suf = pair_suffix - pair_sum; + if (right_suf >= TopK && next_suf < TopK) { + decode_smem[SBASE + sTHR] = 2 * tx + 1; + } + } + __syncthreads(); + + const int threshold = decode_smem[SBASE + sTHR]; + + // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- + int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; + const int sOUT_abs = SBASE + sOUT; + const int sBUF0_abs = SBASE + sBUF0; + + { + const uint32_t uthr = static_cast(threshold); + int item = 0; + const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; + + for (int iter = 0; iter < n_vec_iters; iter++) { + const int i = tx + iter * kThreadsPerBlock; + const bool vec_valid = (i < n_vec); + const int base_idx = i << 2; + +#pragma unroll 4 + for (int sub = 0; sub < 4; sub++) { + const int elem_idx = base_idx + sub; + uint32_t bin = 0; + if (vec_valid) bin = reg_bins[item++]; + const bool is_above = vec_valid && (bin > uthr); + const bool is_equal = vec_valid && (bin == uthr); + + const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); + if (above_mask) { + const int above_count = __popc(above_mask); + const int above_rank = __popc(above_mask & ((1u << lane) - 1)); + int above_base; + if (lane == 0) { + above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); + } + above_base = __shfl_sync(0xffffffff, above_base, 0); + if (is_above) { + output_indices[above_base + above_rank] = elem_idx; + } + } + + const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); + if (equal_mask) { + const int equal_count = __popc(equal_mask); + const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); + int equal_base; + if (lane == 0) { + equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); + } + equal_base = __shfl_sync(0xffffffff, equal_base, 0); + if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { + bufs[0][equal_base + equal_rank] = elem_idx; + } + } + } + } + } + __syncthreads(); + + int remaining_k = TopK - decode_smem[SBASE + sOUT]; + if (remaining_k <= 0) return; + + // If all buffered elements fit, output them all (common for short seqs) + const int raw_buf0 = decode_smem[SBASE + sBUF0]; + if (raw_buf0 <= remaining_k) { + const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; + const int base = decode_smem[SBASE + sOUT]; + for (int i = tx; i < nb; i += kThreadsPerBlock) { + output_indices[base + i] = bufs[0][i]; + } + __syncthreads(); + return; + } + + // ---- Phase 3: Deferred refinement (rare path) ---- + int* refine[2] = {decode_smem, decode_smem + RHIST}; + const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; + + for (int i = tx; i < RHIST; i += kThreadsPerBlock) { + refine[0][i] = 0; + } + __syncthreads(); + + for (int i = tx; i < num_buf0; i += kThreadsPerBlock) { + const uint32_t fp32 = convert_to_uint32_v2(logits[bufs[0][i]]); + atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1); + } + __syncthreads(); + + auto compute_suffix_sum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (tx < RADIX) { + const int stride = 1 << i; + const int s = i & 1; + const int d = s ^ 1; + int value = refine[s][tx]; + if (tx < RADIX - stride) value += refine[s][tx + stride]; + refine[d][tx] = value; + } + __syncthreads(); + } + }; + +#pragma unroll 4 + for (int pass = 0; pass < 4; ++pass) { + const int src = pass & 1; + const int dst = src ^ 1; + + const int raw_buf = decode_smem[SBASE + sBUF0 + src]; + const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; + + compute_suffix_sum(); + + if (tx < RADIX && refine[0][tx] > remaining_k && + refine[0][tx + 1] <= remaining_k) { + decode_smem[SBASE + sREF] = tx; + decode_smem[SBASE + sBUF0 + dst] = 0; + decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; + } + __syncthreads(); + + const int ref_thr = decode_smem[SBASE + sREF]; + remaining_k -= refine[0][ref_thr + 1]; + const int bit_offset = 24 - pass * 8; + + if (remaining_k == 0) { + for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { + const int idx = bufs[src][i]; + const uint32_t fp32 = convert_to_uint32_v2(logits[idx]); + if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { + const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); + output_indices[pos] = idx; + } + } + __syncthreads(); + break; + } + + __syncthreads(); + if (tx < RADIX + 1) refine[0][tx] = 0; + __syncthreads(); + + for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { + const int idx = bufs[src][i]; + const float logit_val = logits[idx]; + const uint32_t fp32 = convert_to_uint32_v2(logit_val); + const int bin = (fp32 >> bit_offset) & 0xFF; + + if (bin > ref_thr) { + const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); + output_indices[pos] = idx; + } else if (bin == ref_thr) { + if (pass == 3) { + const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); + if (slot > 0) output_indices[TopK - slot] = idx; + } else { + const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); + if (__builtin_expect(bp < DBUF, 1)) { + bufs[dst][bp] = idx; + const int nbo = bit_offset - 8; + atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1); + } + } + } + } + __syncthreads(); + } +} + +// ============================================================================ +// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement +// For sequences 8K < seq_len <= 64K. +// ============================================================================ + +// Adapted from: +// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87 +// by: DarkSharpness +// which at the same time is an optimized topk kernel copied from tilelang +// kernel +template +__device__ __noinline__ void histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, + int logits_offset, int seq_len) { + // All shared state lives in dynamic shared memory to avoid static + extern __shared__ char medium_smem[]; + + int (*shared_histogram)[RADIX + 128] = + reinterpret_cast(medium_smem); + int* medium_scalars = reinterpret_cast(medium_smem + kMediumHistBytes); + int& shared_output_count = medium_scalars[0]; + int& shared_threshold_bin = medium_scalars[1]; + int* shared_buffered_count = &medium_scalars[2]; + int& shared_final_k = medium_scalars[4]; + int (*buffered_indices)[MAX_BUFFERED_ITEMS] = + reinterpret_cast(medium_smem + + kMediumHeaderSize); + + const int thread_id = threadIdx.x; + int remaining_k = TopK; + + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const auto bin = convert_to_uint8(logits[idx + logits_offset]); + atomicAdd(&shared_histogram[0][bin], 1); + } + __syncthreads(); + + auto compute_cumulative_sum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (__builtin_expect(thread_id < RADIX, 1)) { + const int stride = 1 << i; + const int src_buffer = i & 1; + const int dst_buffer = src_buffer ^ 1; + int value = shared_histogram[src_buffer][thread_id]; + if (thread_id < RADIX - stride) { + value += shared_histogram[src_buffer][thread_id + stride]; + } + shared_histogram[dst_buffer][thread_id] = value; + } + __syncthreads(); + } + }; + + compute_cumulative_sum(); + + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; + shared_buffered_count[0] = 0; + shared_output_count = 0; + } + __syncthreads(); + + const int threshold_bin = shared_threshold_bin; + remaining_k -= shared_histogram[0][threshold_bin + 1]; + + if (remaining_k == 0) { + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const int bin = convert_to_uint8(logits[idx + logits_offset]); + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + } + __syncthreads(); + return; + } + + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + const int bin = convert_to_uint8(logit_value); + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); + const int next_bin = (fp32_bits >> 24) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + } + __syncthreads(); + +#pragma unroll 4 + for (int pass = 0; pass < 4; ++pass) { + const int src_buffer = pass % 2; + const int dst_buffer = src_buffer ^ 1; + const int raw_buffered = shared_buffered_count[src_buffer]; + const int num_buffered = + (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS; + + compute_cumulative_sum(); + + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; + shared_buffered_count[dst_buffer] = 0; + shared_final_k = remaining_k - shared_histogram[0][thread_id + 1]; + } + __syncthreads(); + + const int threshold_bin = shared_threshold_bin; + remaining_k -= shared_histogram[0][threshold_bin + 1]; + const int bit_offset = 24 - pass * 8; + + if (remaining_k == 0) { + for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { + const int idx = buffered_indices[src_buffer][i]; + const uint32_t fp32_bits = + convert_to_uint32_v2(logits[idx + logits_offset]); + const int bin = (fp32_bits >> bit_offset) & 0xFF; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + } + __syncthreads(); + break; + } + + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { + const int idx = buffered_indices[src_buffer][i]; + const float logit_value = logits[idx + logits_offset]; + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); + const int bin = (fp32_bits >> bit_offset) & 0xFF; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + if (pass == 3) { + const int slot = atomicAdd(&shared_final_k, -1); + if (slot > 0) { + output_indices[TopK - slot] = idx; + } + } else { + const int buffer_pos = + atomicAdd(&shared_buffered_count[dst_buffer], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[dst_buffer][buffer_pos] = idx; + const int next_bit_offset = bit_offset - 8; + const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + } + } + __syncthreads(); + } +} + +// ============================================================================ +// Inter-CTA sync primitives +// ============================================================================ + +__device__ __forceinline__ int ld_acquire(int* ptr) { + int state = 0; +#if (__CUDA_ARCH__ >= 700) + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(ptr)); +#else + asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); +#endif + return state; +} + +__device__ __forceinline__ void red_release(int* ptr, int val) { +#if (__CUDA_ARCH__ >= 700) + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(ptr), "r"(val)); +#else + __threadfence(); + atomicAdd(ptr, val); +#endif +} + +__device__ __forceinline__ void st_release(int* ptr, int val) { +#if (__CUDA_ARCH__ >= 700) + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val)); +#else + __threadfence(); + atomicExch(ptr, val); +#endif +} + +__device__ __forceinline__ void wait_ge(int* ptr, int target_val, + int thread_idx) { + if (thread_idx == 0) { +#pragma unroll 1 + while (ld_acquire(ptr) < target_val) { + } + } + __syncthreads(); +} + +// ============================================================================ +// Large path: multi-CTA radix select for sequences > 64K +// +// Each row is processed by a group of CTAs. Each CTA loads its chunk into +// shared memory as ordered uint32, then participates in 4 rounds of +// coordinated radix select via global-memory histograms and barriers. +// ============================================================================ + +// ============================================================================ +// Multi-CTA cooperative RadixTopK for a single large row. +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 +// ============================================================================ + +template +__device__ void radix_topk(const float* __restrict__ row_input, + int32_t* __restrict__ row_output, uint32_t seq_len, + uint32_t my_chunk_start, uint32_t chunk_size, + uint32_t* local_histogram, uint32_t* suffix_sum, + uint32_t* shared_scalars, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t iter, uint32_t tx) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; + const uint32_t actual_chunk_size = + (my_chunk_start < seq_len) ? (my_chunk_end - my_chunk_start) : 0; + + // -- Stage 1: Load chunk to shared memory as ordered uint32 -- + { + const uint32_t aligned_size = (actual_chunk_size / VEC_SIZE) * VEC_SIZE; + + for (uint32_t i = tx * VEC_SIZE; i < aligned_size; + i += kThreadsPerBlock * VEC_SIZE) { + const float* src = row_input + my_chunk_start + i; + if constexpr (VEC_SIZE == 4) { + float4 v = *reinterpret_cast(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); + shared_ordered[i + 2] = convert_to_uint32_v2(v.z); + shared_ordered[i + 3] = convert_to_uint32_v2(v.w); + } else if constexpr (VEC_SIZE == 2) { + float2 v = *reinterpret_cast(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); + } else { + shared_ordered[i] = convert_to_uint32_v2(*src); + } + } + for (uint32_t i = aligned_size + tx; i < actual_chunk_size; + i += kThreadsPerBlock) { + shared_ordered[i] = convert_to_uint32_v2(row_input[my_chunk_start + i]); + } + } + __syncthreads(); + + // -- Init radix select state -- + if (tx == 0) { + shared_scalars[0] = 0; // prefix + shared_scalars[1] = TopK; // remaining_k + } + __syncthreads(); + + // -- Initial barrier -- + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + if (cta_in_group == 0 && tx == 0) { + st_release(&state->output_counter, 0); + } + + // -- Stage 2: 4 rounds of radix select -- + for (uint32_t round = 0; round < 4; round++) { + const uint32_t global_round = iter * 4 + round; + const uint32_t shift = 24 - round * 8; + const uint32_t prefix = shared_scalars[0]; + const uint32_t remaining_k = shared_scalars[1]; + + uint32_t* current_hist = state->histogram[global_round % 3]; + uint32_t* next_hist = state->histogram[(global_round + 1) % 3]; + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + local_histogram[i] = 0; + } + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + uint32_t ordered = shared_ordered[i]; + uint32_t mask = (round == 0) ? 0u : (~0u << (32 - round * 8)); + if ((ordered & mask) == prefix) { + uint32_t bucket = (ordered >> shift) & 0xFF; + atomicAdd(&local_histogram[bucket], 1); + } + } + __syncthreads(); + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + if (local_histogram[i] > 0) { + atomicAdd(¤t_hist[i], local_histogram[i]); + } + } + + if (cta_in_group == 0) { + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + next_hist[i] = 0; + } + } + + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + suffix_sum[i] = current_hist[i]; + } + __syncthreads(); + + for (uint32_t stride = 1; stride < RADIX; stride *= 2) { + uint32_t val = 0; + if (tx < RADIX) { + val = suffix_sum[tx]; + if (tx + stride < RADIX) val += suffix_sum[tx + stride]; + } + __syncthreads(); + if (tx < RADIX) suffix_sum[tx] = val; + __syncthreads(); + } + + if (tx == 0) { + shared_scalars[2] = 0; + shared_scalars[3] = remaining_k; + } + __syncthreads(); + + if (tx < RADIX) { + uint32_t count_ge = suffix_sum[tx]; + uint32_t count_gt = (tx + 1 < RADIX) ? suffix_sum[tx + 1] : 0; + if (count_ge >= remaining_k && count_gt < remaining_k) { + shared_scalars[2] = tx; + shared_scalars[3] = remaining_k - count_gt; + } + } + __syncthreads(); + + if (tx == 0) { + shared_scalars[0] = prefix | (shared_scalars[2] << shift); + shared_scalars[1] = shared_scalars[3]; + } + __syncthreads(); + } // end 4 radix rounds + + // -- Count local > pivot elements -- + const uint32_t ordered_pivot = shared_scalars[0]; + + if (tx == 0) suffix_sum[0] = 0; + __syncthreads(); + + uint32_t my_gt_count = 0; + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] > ordered_pivot) my_gt_count++; + } + for (int offset = 16; offset > 0; offset /= 2) { + my_gt_count += __shfl_down_sync(0xffffffff, my_gt_count, offset); + } + if (tx % 32 == 0 && my_gt_count > 0) { + atomicAdd(&suffix_sum[0], my_gt_count); + } + __syncthreads(); + const uint32_t local_gt_count = suffix_sum[0]; + + // -- Stage 3: Collect top-k indices -- + if (tx == 0) { + local_histogram[0] = 0; + if (local_gt_count > 0) { + local_histogram[1] = + atomicAdd(&state->output_counter, static_cast(local_gt_count)); + } + } + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] > ordered_pivot) { + uint32_t local_pos = atomicAdd(&local_histogram[0], 1); + int pos = static_cast(local_histogram[1]) + local_pos; + row_output[pos] = static_cast(my_chunk_start + i); + } + } + + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] == ordered_pivot) { + int pos = atomicAdd(&state->output_counter, 1); + if (pos < TopK) { + row_output[pos] = static_cast(my_chunk_start + i); + } + } + } +} + +// ============================================================================ +// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK +// BS>32 uses standalone histogram_256_buffered_topk (separate kernel, +// see filtered_topk.cuh) +// ============================================================================ + +template +__global__ void __launch_bounds__(kThreadsPerBlock, 2) + persistent_topk_kernel(PersistentTopKParams params) { + const uint32_t tx = threadIdx.x; + extern __shared__ uint8_t smem_raw[]; + + // ======================================================================== + // Group mode: multi-CTA groups with static round-robin row assignment. + // Non-large rows: CTA-0 handles trivial/decode/medium. + // Large rows: all CTAs in the group cooperate via RadixTopK. + // ======================================================================== + const uint32_t ctas_per_group = params.ctas_per_group; + const uint32_t group_id = blockIdx.x / ctas_per_group; + const uint32_t cta_in_group = blockIdx.x % ctas_per_group; + const uint32_t num_groups = gridDim.x / ctas_per_group; + const uint32_t chunk_size = params.chunk_size; + + if (blockIdx.x >= num_groups * ctas_per_group) return; + + // Early exit: non-CTA-0 threads are never needed if no large rows exist + if (cta_in_group != 0 && params.max_seq_len <= RADIX_THRESHOLD) return; + + uint32_t* local_histogram = reinterpret_cast(smem_raw); + uint32_t* suffix_sum = local_histogram + RADIX; + uint32_t* shared_scalars = suffix_sum + RADIX; + uint32_t* shared_ordered = + reinterpret_cast(smem_raw + kFixedSmemLarge); + + // RadixRowState for multi-CTA cooperative radix. + // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu + // before launch — that gives a stream-ordered happens-before edge for all + // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA + // __syncthreads) did not provide and which manifested as a race against + // CTA-1+'s first red_release on arrival_counter. + RadixRowState* state = ¶ms.row_states[group_id]; + + int barrier_phase = 0; + const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups; + + for (uint32_t iter = 0; iter < total_iters; iter++) { + // Static round-robin: all CTAs in the group implicitly agree on the row + uint32_t row_idx = group_id + iter * num_groups; + if (row_idx >= params.num_rows) break; + + const uint32_t seq_len = params.lengths[row_idx]; + int32_t* row_output = params.output + row_idx * params.top_k; + const float* row_input = params.input + row_idx * params.stride; + + if (seq_len <= RADIX_THRESHOLD) { + if (cta_in_group == 0) { + if (seq_len <= static_cast(TopK)) { + // Trivial case: seq_len <= TopK + for (uint32_t i = tx; i < static_cast(TopK); + i += kThreadsPerBlock) { + row_output[i] = (i < seq_len) ? static_cast(i) : -1; + } + } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { + histogram_2048_topk(row_input, row_output, seq_len); + } else { + histogram_256_topk(row_input, row_output, 0, seq_len); + } + } + continue; + } + + const uint32_t my_chunk_start = cta_in_group * chunk_size; + radix_topk( + row_input, row_output, seq_len, my_chunk_start, chunk_size, + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, iter, tx); + } +} + +} // namespace persistent + +// ============================================================================ +// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only. +// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent). +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 +// ============================================================================ + +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + return e; \ + } \ + } + +#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ + +template +struct vec_t { + T data[N]; + + FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; } + FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; } + + FLASHINFER_INLINE void cast_load(const T* ptr) { +#pragma unroll + for (size_t i = 0; i < N; ++i) { + data[i] = ptr[i]; + } + } + + FLASHINFER_INLINE void cast_store(T* ptr) const { +#pragma unroll + for (size_t i = 0; i < N; ++i) { + ptr[i] = data[i]; + } + } +}; +#undef FLASHINFER_INLINE + +// FilteredTopK traits for different data types +template +struct FilteredTopKTraits; + +// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4 +// refinement rounds +template <> +struct FilteredTopKTraits { + using OrderedType = uint32_t; + static constexpr int NUM_REFINE_ROUNDS = 4; + static constexpr int FIRST_REFINE_SHIFT = 24; + + __device__ __forceinline__ static uint8_t ToCoarseKey(float x) { + // Convert to FP16 representation and extract high 8 bits + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return static_cast(key >> 8); + } + + __device__ __forceinline__ static OrderedType ToOrdered(float x) { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + } +}; + +constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024; +constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = + 16 * 1024; // 16K indices per buffer +constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = + sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; + constexpr int RADIX = 256; + constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; + + const uint32_t bid = blockIdx.x; + const int tx = threadIdx.x; + + if (bid >= num_rows) return; + + const int length = + (lengths != nullptr) ? lengths[bid] : static_cast(max_len); + const DType* score = input + bid * max_len; + IdType* dst = output + bid * top_k; + + // Trivial case: length <= top_k + if (length <= static_cast(top_k)) { + for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { + dst[i] = (i < length) ? static_cast(i) : static_cast(-1); + } + return; + } + + // Static shared memory + alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; + alignas(128) __shared__ int s_counter; + alignas(128) __shared__ int s_threshold_bin_id; + alignas(128) __shared__ int s_num_input[2]; + alignas(128) __shared__ int s_indices[MAX_K]; + + auto& s_histogram = s_histogram_buf[0]; + + // Dynamic shared memory for input double buffer + extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE]; + + using Traits = FilteredTopKTraits; + int topk = top_k; + + // Stage 1: 8-bit coarse histogram with vectorized loads + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + + vec_t score_vec; + + const int aligned_length = (length / VEC_SIZE) * VEC_SIZE; +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const auto bin = Traits::ToCoarseKey(score_vec[j]); + atomicAdd(&s_histogram[bin], 1); + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + const auto bin = Traits::ToCoarseKey(score[i]); + atomicAdd(&s_histogram[bin], 1); + } + __syncthreads(); + + // Suffix sum + const auto run_cumsum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (tx < RADIX) { + const auto j = 1 << i; + const auto k = i & 1; + auto value = s_histogram_buf[k][tx]; + if (tx < RADIX - j) { + value += s_histogram_buf[k][tx + j]; + } + s_histogram_buf[k ^ 1][tx] = value; + } + __syncthreads(); + } + }; + + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + s_num_input[0] = 0; + s_counter = 0; + } + __syncthreads(); + + const auto threshold_bin = s_threshold_bin_id; + topk -= s_histogram[threshold_bin + 1]; + + constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; + constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; + + if (topk == 0) { + // Collect indices where bin > threshold +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const auto bin = static_cast(Traits::ToCoarseKey(score_vec[j])); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = base + j; + } + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + const auto bin = static_cast(Traits::ToCoarseKey(score[i])); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = i; + } + } + __syncthreads(); + } else { + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + + // Filter + histogram for refinement + auto filter_and_add_to_histogram = [&](auto raw_input, int index) { + const auto bin = static_cast(Traits::ToCoarseKey(raw_input)); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = index; + } else if (bin == threshold_bin) { + const auto pos = atomicAdd(&s_num_input[0], 1); + if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { + s_input_idx[0][pos] = index; + const auto ordered = Traits::ToOrdered(raw_input); + const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; + atomicAdd(&s_histogram[sub_bin], 1); + } + } + }; +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + filter_and_add_to_histogram(score_vec[j], base + j); + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + filter_and_add_to_histogram(score[i], i); + } + __syncthreads(); + + // Stage 2: refine with 8bit radix passes +#pragma unroll + for (int round = 0; round < NUM_ROUNDS; ++round) { + __shared__ int s_last_remain; + const auto r_idx = round % 2; + + const auto _raw_num_input = s_num_input[r_idx]; + const auto num_input = + (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; + + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + s_num_input[r_idx ^ 1] = 0; + s_last_remain = topk - s_histogram[tx + 1]; + } + __syncthreads(); + + const auto threshold = s_threshold_bin_id; + topk -= s_histogram[threshold + 1]; + + const int offset = FIRST_SHIFT - round * 8; + const bool is_last_round = (round == NUM_ROUNDS - 1); + + if (topk == 0) { + for (int i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; + if (static_cast(bin) > threshold) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = idx; + } + } + __syncthreads(); + break; + } else { + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + for (int i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto raw_input = score[idx]; + const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; + if (static_cast(bin) > threshold) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = idx; + } else if (static_cast(bin) == threshold) { + if (is_last_round) { + const auto pos = atomicAdd(&s_last_remain, -1); + if (pos > 0) { + s_indices[top_k - pos] = idx; + } + } else { + const auto pos = atomicAdd(&s_num_input[r_idx ^ 1], 1); + if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { + s_input_idx[r_idx ^ 1][pos] = idx; + const auto bin32 = Traits::ToOrdered(raw_input); + const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; + atomicAdd(&s_histogram[sub_bin], 1); + } + } + } + } + __syncthreads(); + } + } + } + + // Output phase - mode-specific +#pragma unroll 2 + for (int base = tx; base < static_cast(top_k); base += BLOCK_SIZE) { + const int idx = s_indices[base]; + dst[base] = static_cast(idx); + } +} + +// Helper to compute GCD for VEC_SIZE selection +constexpr uint32_t gcd(uint32_t a, uint32_t b) { + while (b != 0) { + uint32_t t = b; + b = a % b; + a = t; + } + return a; +} + +// Compute optimal VEC_SIZE based on max_len and dtype +// Returns 1, 2, 4, or 8 +template +constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { + constexpr int MAX_VEC = 16 / sizeof(DType); // 4 for float32, 8 for fp16/bf16 + // Use GCD to find largest power-of-2 divisor + const uint32_t g = gcd(max_len, static_cast(MAX_VEC)); + return static_cast(g); +} + +template +cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, + IdType* lengths, uint32_t num_rows, + uint32_t top_k_val, uint32_t max_len, + cudaStream_t stream = 0) { + constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; + constexpr int MAX_VEC = 16 / sizeof(DType); + + dim3 grid(num_rows); + dim3 block(FILTERED_TOPK_BLOCK_THREADS); + void* args[] = {&input, &output_indices, &lengths, + &num_rows, &top_k_val, &max_len}; + + const int vec_size = ComputeFilteredTopKVecSize(max_len); + +#define DISPATCH_VEC_SIZE(VS) \ + if (vec_size == VS) { \ + auto kernel = FilteredTopKUnifiedKernel; \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ + smem_size, stream)); \ + return cudaSuccess; \ + } + + DISPATCH_VEC_SIZE(1) + DISPATCH_VEC_SIZE(2) + DISPATCH_VEC_SIZE(4) + if constexpr (MAX_VEC >= 8) { + DISPATCH_VEC_SIZE(8) + } +#undef DISPATCH_VEC_SIZE + + return cudaSuccess; +} + +} // namespace vllm + +#endif // PERSISTENT_TOPK_CUH_ diff --git a/csrc/pos_encoding_kernels.cu b/csrc/pos_encoding_kernels.cu index b5645b33b907..c45ebd34729b 100644 --- a/csrc/pos_encoding_kernels.cu +++ b/csrc/pos_encoding_kernels.cu @@ -9,28 +9,29 @@ namespace vllm { template inline __device__ void apply_token_rotary_embedding( - scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr, - const scalar_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; - scalar_t cos, sin; + float cos_f, sin_f; if (IS_NEOX) { - // GPT-NeoX style rotary embedding. x_index = rot_offset; y_index = embed_dim + rot_offset; - cos = VLLM_LDG(cos_ptr + x_index); - sin = VLLM_LDG(sin_ptr + x_index); + cos_f = VLLM_LDG(cos_ptr + x_index); + sin_f = VLLM_LDG(sin_ptr + x_index); } else { - // GPT-J style rotary embedding. x_index = 2 * rot_offset; y_index = 2 * rot_offset + 1; - cos = VLLM_LDG(cos_ptr + x_index / 2); - sin = 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); } - - const scalar_t x = arr[x_index]; - const scalar_t y = arr[y_index]; - arr[x_index] = x * cos - y * sin; - arr[y_index] = y * cos + x * sin; + if (inverse) { + sin_f = -sin_f; + } + const float x_f = static_cast(arr[x_index]); + const float y_f = static_cast(arr[y_index]); + arr[x_index] = static_cast(x_f * cos_f - y_f * sin_f); + arr[y_index] = static_cast(y_f * cos_f + x_f * sin_f); } template @@ -42,22 +43,23 @@ inline __device__ void apply_rotary_embedding( // [batch_size, seq_len, num_kv_heads, // head_size] or [num_tokens, num_kv_heads, // head_size] - const scalar_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 head_stride, const int64_t rope_dim_offset, + const bool inverse) { const int embed_dim = rot_dim / 2; - const scalar_t* cos_ptr = cache_ptr; - const scalar_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) { const int head_idx = i / embed_dim; const int64_t token_head = - token_idx * query_stride + head_idx * head_stride; + token_idx * query_stride + head_idx * head_stride + rope_dim_offset; const int rot_offset = i % embed_dim; apply_token_rotary_embedding( - query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); + query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse); } if (key != nullptr) { @@ -65,10 +67,10 @@ inline __device__ void apply_rotary_embedding( for (int i = threadIdx.x; i < nk; i += blockDim.x) { const int head_idx = i / embed_dim; const int64_t token_head = - token_idx * key_stride + head_idx * head_stride; + token_idx * key_stride + head_idx * head_stride + rope_dim_offset; const int rot_offset = i % embed_dim; apply_token_rotary_embedding( - key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); + key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse); } } } @@ -84,19 +86,18 @@ __global__ void rotary_embedding_kernel( // [batch_size, seq_len, num_kv_heads, // head_size] or [num_tokens, num_kv_heads, // head_size] - const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // - // 2] + 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) { - // Each thread block is responsible for one token. + 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 scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; + const float* cache_ptr = cos_sin_cache + pos * rot_dim; apply_rotary_embedding( query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim, - token_idx, query_stride, key_stride, head_stride); + token_idx, query_stride, key_stride, head_stride, rope_dim_offset, + inverse); } } // namespace vllm @@ -115,7 +116,7 @@ void rotary_embedding( // [num_tokens, num_heads, head_size] int64_t head_size, torch::Tensor& cos_sin_cache, // [max_position, rot_dim] - bool is_neox) { + bool is_neox, int64_t rope_dim_offset, bool inverse) { // num_tokens = batch_size * seq_len int64_t num_tokens = positions.numel(); int positions_ndim = positions.dim(); @@ -154,6 +155,8 @@ void rotary_embedding( int seq_dim_idx = positions_ndim - 1; int64_t query_stride = query.stride(seq_dim_idx); int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0; + + TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size); // Determine head stride: for [*, heads, head_size] use stride of last dim; // for flat [*, heads*head_size], heads blocks are contiguous of size // head_size @@ -165,20 +168,23 @@ void rotary_embedding( dim3 block(std::min(num_heads * rot_dim / 2, 512)); const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto cache_f32 = cos_sin_cache.to(torch::kFloat32); VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] { if (is_neox) { vllm::rotary_embedding_kernel<<>>( positions.data_ptr(), query.data_ptr(), key.has_value() ? key->data_ptr() : nullptr, - cos_sin_cache.data_ptr(), rot_dim, query_stride, key_stride, - head_stride, num_heads, num_kv_heads, head_size); + cache_f32.data_ptr(), rot_dim, query_stride, key_stride, + head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset, + inverse); } else { vllm::rotary_embedding_kernel <<>>( positions.data_ptr(), query.data_ptr(), key.has_value() ? key->data_ptr() : nullptr, - cos_sin_cache.data_ptr(), rot_dim, query_stride, - key_stride, head_stride, num_heads, num_kv_heads, head_size); + cache_f32.data_ptr(), rot_dim, query_stride, key_stride, + head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset, + inverse); } }); } diff --git a/csrc/quantization/fused_kernels/quant_conversions.cuh b/csrc/quantization/fused_kernels/quant_conversions.cuh index 3711c47edc8c..fc60643777e0 100644 --- a/csrc/quantization/fused_kernels/quant_conversions.cuh +++ b/csrc/quantization/fused_kernels/quant_conversions.cuh @@ -6,7 +6,7 @@ #include "libtorch_stable/quantization/vectorization.cuh" // TODO(luka/varun):refactor common.cuh to use this file instead -#include "quantization/w8a8/fp8/common.cuh" +#include "../w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 81f5cb83f3e1..7ae644d81d48 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -639,7 +639,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { // function with template. #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ - if (KV_DTYPE == "auto") { \ + vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ + vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ + if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ if (SRC_DTYPE == at::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ } else if (SRC_DTYPE == at::ScalarType::Half) { \ @@ -649,21 +651,18 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { } else { \ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ - } else { \ - if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/quantization/w8a8/fp8/common.cuh b/csrc/quantization/w8a8/fp8/common.cuh index 7a385f5163ae..7576f7179501 100644 --- a/csrc/quantization/w8a8/fp8/common.cuh +++ b/csrc/quantization/w8a8/fp8/common.cuh @@ -1,7 +1,7 @@ #pragma once #include "libtorch_stable/quantization/vectorization.cuh" -#include "quantization/utils.cuh" +#include "../../utils.cuh" #include diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 421e8092474b..3b7e25dc56be 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -543,7 +543,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { // function with template. #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ - if (KV_DTYPE == "auto") { \ + vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ + vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ + if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ if (SRC_DTYPE == at::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ } else if (SRC_DTYPE == at::ScalarType::Half) { \ @@ -553,43 +555,28 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { } else { \ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ - } else { \ - if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ - } else if (KV_DTYPE == "fp8_e5m2") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ - } else if (KV_DTYPE == "fp8_ds_mla") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else { \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index a339c5641bb4..9e6c0726d19e 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -40,15 +40,6 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz; #define __HIP__FP8MFMA__ #endif -#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \ - defined(__gfx1150__) || defined(__gfx1151__)) - #define __HIP__GFX11__ -#endif - -#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__)) - #define __HIP__GFX12__ -#endif - #if defined(NDEBUG) #undef NDEBUG #include @@ -1629,7 +1620,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( } } -#elif defined(__HIP__GFX11__) +#elif defined(__GFX11__) using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; @@ -2388,7 +2379,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( out_ptr[threadIdx.x] = from_float(acc); } -#elif defined(__HIP__GFX12__) +#elif defined(__GFX12__) using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 60e10e53391a..3342db37be90 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -26,16 +26,11 @@ #define __HIP__GFX9__ #endif -#if defined(__HIPCC__) && \ - (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \ - defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__)) +// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts +#if defined(__GFX11__) || defined(__GFX12__) #define __HIP__GFX1X__ #endif -#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__)) - #define __HIP__GFX12__ -#endif - #if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) #define __HIP__MI3XX__ #endif @@ -1845,7 +1840,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b, return out_c; } -#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +#if defined(__HIP__MI3XX__) || defined(__GFX12__) template __global__ void __launch_bounds__(WvPrGrp* THRDS) @@ -1893,7 +1888,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) float sB = *s_B; while (m < M) { - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8 float sum[N][YTILE] = {}; #else @@ -1931,7 +1926,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { for (uint32_t n = 0; n < N; n++) { - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4) for (int y = 0; y < YTILE; ++y) { #pragma unroll @@ -1955,7 +1950,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } // Final reduction - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { @@ -1993,7 +1988,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #endif const bool writeback_lane = - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ threadIdx.x == (THRDS - 1); #else threadIdx.x == 0; @@ -2009,7 +2004,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { if (y + m >= M) break; // To avoid mem access fault. - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ float result = sum[n][y] * sA * sB; #else float result = sum[n][y][0] * sA * sB; @@ -2027,7 +2022,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) m += CuCount * _WvPrGrp * YTILE; } } -#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__) +#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__) template __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, @@ -2039,9 +2034,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int _WvPrGrp, const int CuCount) { UNREACHABLE_CODE } -#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +#endif // defined(__HIP__MI3XX__) || defined(__GFX12__) -#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +#if defined(__HIP__MI3XX__) || defined(__GFX12__) template __global__ void __launch_bounds__(WvPrGrp* THRDS) @@ -2088,7 +2083,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) float sB = *s_B; while (m < M) { - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8 float sum[N][YTILE] = {}; #else @@ -2128,7 +2123,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { for (uint32_t n = 0; n < N; n++) { - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4) for (int y = 0; y < YTILE; ++y) { #pragma unroll @@ -2152,7 +2147,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } // Final reduction - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { @@ -2190,7 +2185,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #endif const bool writeback_lane = - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ threadIdx.x == (THRDS - 1); #else threadIdx.x == 0; @@ -2206,7 +2201,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (int n = 0; n < N; n++) { for (int y = 0; y < YTILE; y++) { if (y + m >= M) break; // To avoid mem access fault. - #ifdef __HIP__GFX12__ + #ifdef __GFX12__ float result = sum[n][y] * sA * sB; #else float result = sum[n][y][0] * sA * sB; @@ -2224,7 +2219,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) m += CuCount * _WvPrGrp * YTILE; } } -#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__) +#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__) template __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, @@ -2236,7 +2231,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int CuCount) { UNREACHABLE_CODE } -#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +#endif // defined(__HIP__MI3XX__) || defined(__GFX12__) void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a, const std::optional& in_bias, at::Tensor& out_c, diff --git a/csrc/sampler.cu b/csrc/sampler.cu index 2e76873c8f18..14d84013c08d 100644 --- a/csrc/sampler.cu +++ b/csrc/sampler.cu @@ -258,7 +258,13 @@ __device__ bool processHistogramStep( auto processBins = [&](float logit, int idx) { if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - if (binIdx < thresholdBinIdx) { + // Only write elements with binIdx < thresholdBinIdx when: + // 1. This is step 0 and the threshold bin is small enough (no step 1) + // 2. This is step >= 1 (where pattern matching filters correctly) + // This prevents duplicates when step 0 and step 1 both run. + bool shouldWriteDirectly = + (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1); + if (binIdx < thresholdBinIdx && shouldWriteDirectly) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -564,8 +570,9 @@ template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( const float* logits, const int* seqLens, int* outIndices, int stride0, - int stride1, const int topK, int next_n, float* outLogits = nullptr, - const int numBlocksToMerge = 0, const int* indices = nullptr) { + int stride1, const int topK, int next_n, int seqLensIs2D = 0, + float* outLogits = nullptr, const int numBlocksToMerge = 0, + const int* indices = nullptr) { // The number of bins in the histogram. static constexpr int kNumBins = 2048; @@ -574,8 +581,16 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( // The range of logits within the row. int rowStart = 0; - int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1); + int batch_idx = rowIdx / next_n; + int next_n_idx = rowIdx % next_n; + // seqLensIs2D=0: 1D seqLens — all rows in a batch share the same seq_len; + // kernel computes per-row effective length via offset. + // seqLensIs2D=1: 2D seqLens — each logit row has its own pre-computed + // effective length (flat index rowIdx = b*next_n + j maps + // directly to seqLens[b, j] in C-contiguous layout). + int seq_len = seqLensIs2D ? seqLens[rowIdx] : seqLens[batch_idx]; + int rowEnd = + seqLensIs2D ? max(0, seq_len) : max(0, seq_len - next_n + next_n_idx + 1); // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) { @@ -653,6 +668,11 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const auto numColumns = logits.size(1); + // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed + // effective seq_len. False if seqLens is 1D (B,): all rows in a batch share + // the same seq_len and the kernel computes the per-row offset itself. + int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode @@ -660,7 +680,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n)); + static_cast(next_n), seqLensIs2D); } else if (numColumns < kSplitWorkThreshold) { // From this threshold, use radix sort instead vllm::topKPerRowDecode @@ -668,7 +688,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n)); + static_cast(next_n), seqLensIs2D); } else { // Long sequences are run in two steps constexpr auto multipleBlocksPerRowConfig = 10; @@ -686,15 +706,16 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), outIndicesAux.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n), outLogitsAux.data_ptr()); + static_cast(next_n), seqLensIs2D, + outLogitsAux.data_ptr()); constexpr int kNumThreadsPerBlockMerge = 1024; vllm::topKPerRowDecode <<>>( outLogitsAux.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), multipleBlocksPerRowConfig * topK, 1, - static_cast(topK), static_cast(next_n), nullptr, - multipleBlocksPerRowConfig, outIndicesAux.data_ptr()); + static_cast(topK), static_cast(next_n), seqLensIs2D, + nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr()); } } diff --git a/csrc/topk.cu b/csrc/topk.cu index a7850f5363b9..68352629ef02 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -1,373 +1,258 @@ -// Portions of this file are adapted from SGLang PR: -// https://github.com/sgl-project/sglang/pull/11194 -// and -// https://github.com/sgl-project/sglang/pull/17747 +// Persistent TopK kernel for DeepSeek V3 sparse attention indexer. +// See persistent_topk.cuh for kernel implementation. -#include "cuda_compat.h" -#include "dispatch_utils.h" - -#include -#include +#include +#include +#include +#include #ifndef USE_ROCM - #include -#else - #include + #include "persistent_topk.cuh" #endif -namespace vllm { - -constexpr int TopK = 2048; // DeepSeek V3 sparse attention top-k -constexpr int kThreadsPerBlock = 1024; // Threads per block - -// Shared memory budget -#if defined(USE_ROCM) -constexpr size_t kSmem = 48 * 1024; // ROCm default: 48KB -#else -// Reduced from 128KB to 32KB to improve occupancy. -// Each radix pass needs at most ~TopK candidates in the threshold bin, -// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient. -constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes) -#endif - -struct FastTopKParams { - const float* __restrict__ input; // [batch, seq_len] Logits - const int32_t* __restrict__ row_starts; // [batch] Offset into each row - // (optional) - int32_t* __restrict__ indices; // [batch, TopK] Output top-k indices - int32_t* __restrict__ lengths; // [batch] Sequence lengths per row - int64_t input_stride; // Stride between rows -}; - -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); -} - -__device__ void naive_topk_cuda(const float* __restrict__ logits, - int32_t* __restrict__ output_indices, - int32_t seq_len) { - const int thread_id = threadIdx.x; - for (int i = thread_id; i < TopK; i += kThreadsPerBlock) { - output_indices[i] = (i < seq_len) ? i : -1; - } -} - -// Adapted from: -// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87 -// by: DarkSharpness -// which at the same time is an optimized topk kernel copied from tilelang -// kernel -__device__ void fast_topk_cuda_tl( - const float* __restrict__ logits, // Input logits [seq_len] - int* __restrict__ output_indices, // Output top-k indices [TopK] - int logits_offset, // Starting offset in logits array - int seq_len) // Number of valid logits to process -{ - constexpr int RADIX = 256; - constexpr int MAX_BUFFERED_ITEMS = kSmem / (2 * sizeof(int)); - - alignas(128) __shared__ int shared_histogram[2][RADIX + 128]; - alignas(128) __shared__ int shared_output_count; - alignas(128) __shared__ int shared_threshold_bin; - alignas(128) __shared__ int shared_buffered_count[2]; - - extern __shared__ int buffered_indices[][MAX_BUFFERED_ITEMS]; +namespace { - const int thread_id = threadIdx.x; - int remaining_k = TopK; - - // Pass 0: Build coarse 8-bit histogram using FP16 high bits - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); - - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const auto bin = convert_to_uint8(logits[idx + logits_offset]); - ::atomicAdd(&shared_histogram[0][bin], 1); +#ifndef USE_ROCM +template +void launch_persistent_topk(const torch::Tensor& logits, + const torch::Tensor& lengths, torch::Tensor& output, + torch::Tensor& workspace, int64_t max_seq_len) { + namespace P = vllm::persistent; + + const int64_t num_rows = logits.size(0); + const int64_t stride = logits.size(1); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + static int num_sms = 0; + static int max_smem_per_block = 0; + if (num_sms == 0) { + int device; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device); + cudaDeviceGetAttribute(&max_smem_per_block, + cudaDevAttrMaxSharedMemoryPerBlockOptin, device); } - __syncthreads(); - - // Helper: Compute cumulative sum (suffix sum) over histogram using ping-pong - // buffers - auto compute_cumulative_sum = [&]() { - static_assert(1 << 8 == RADIX, - "Radix must be 256 for 8 unrolled iterations"); -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (C10_LIKELY(thread_id < RADIX)) { - const int stride = 1 << i; - const int src_buffer = i & 1; - const int dst_buffer = src_buffer ^ 1; - int value = shared_histogram[src_buffer][thread_id]; - if (thread_id < RADIX - stride) { - value += shared_histogram[src_buffer][thread_id + stride]; - } - shared_histogram[dst_buffer][thread_id] = value; - } - __syncthreads(); + if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + cudaError_t status = + vllm::FilteredTopKRaggedTransform( + logits.data_ptr(), output.data_ptr(), + lengths.data_ptr(), static_cast(num_rows), + static_cast(TopK), static_cast(stride), stream); + TORCH_CHECK(status == cudaSuccess, + "FilteredTopK failed: ", cudaGetErrorString(status)); + } else { + TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8"); + + int effective_max_smem; + if (num_rows <= 4) { + effective_max_smem = + std::min(max_smem_per_block, static_cast(P::kSmemMedium)); + } else if (num_rows <= 8) { + constexpr int kSmemCapMedium = 48 * 1024; + effective_max_smem = std::min(max_smem_per_block, kSmemCapMedium); + } else { + effective_max_smem = max_smem_per_block; } - }; - compute_cumulative_sum(); - - // Find threshold bin where cumsum crosses remaining_k - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[0] = 0; - shared_output_count = 0; - } - __syncthreads(); - - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - - // Early exit if threshold bin perfectly matches remaining_k - if (remaining_k == 0) { - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const int bin = convert_to_uint8(logits[idx + logits_offset]); - if (bin > threshold_bin) { - const int output_pos = ::atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } + size_t available_for_ordered = + static_cast(effective_max_smem) - P::kFixedSmemLarge; + uint32_t max_chunk_elements = + static_cast(available_for_ordered / sizeof(uint32_t)); + + uint32_t vec_size = 1; + if (stride % 4 == 0) + vec_size = 4; + else if (stride % 2 == 0) + vec_size = 2; + + max_chunk_elements = (max_chunk_elements / vec_size) * vec_size; + uint32_t min_chunk = vec_size * P::kThreadsPerBlock; + if (max_chunk_elements < min_chunk) max_chunk_elements = min_chunk; + + uint32_t ctas_per_group = + (static_cast(stride) + max_chunk_elements - 1) / + max_chunk_elements; + uint32_t chunk_size = + (static_cast(stride) + ctas_per_group - 1) / ctas_per_group; + chunk_size = ((chunk_size + vec_size - 1) / vec_size) * vec_size; + if (chunk_size > max_chunk_elements) chunk_size = max_chunk_elements; + + 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, P::kThreadsPerBlock, + smem_size); + } else if (vec_size == 2) { + occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, + smem_size); + } else { + occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, + smem_size); } - __syncthreads(); - return; - } - - // Prepare for refinement passes: Process threshold bin - __syncthreads(); - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); - - // Scan all elements and: - // 1. Write indices > threshold_bin to output - // 2. Buffer indices == threshold_bin for refinement - // 3. Build histogram for next refinement pass (fused optimization) - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const float logit_value = logits[idx + logits_offset]; - const int bin = convert_to_uint8(logit_value); - - if (bin > threshold_bin) { - // in top-k, write to output - const int output_pos = ::atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - // Candidate for top-k, needs refinement - const int buffer_pos = ::atomicAdd(&shared_buffered_count[0], 1); - if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) { - buffered_indices[0][buffer_pos] = idx; - // Fused: Build histogram for next pass - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int next_bin = (fp32_bits >> 24) & 0xFF; - ::atomicAdd(&shared_histogram[0][next_bin], 1); + TORCH_CHECK(occ_err == cudaSuccess, + "persistent_topk occupancy query failed: ", + cudaGetErrorString(occ_err)); + 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(max_seq_len) > P::RADIX_THRESHOLD; + + const uint32_t hw_resident_cap = + static_cast(num_sms) * static_cast(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(num_sms) : 1u; + if (max_resident_ctas >= headroom + ctas_per_group) { + max_resident_ctas -= headroom; } } - } - __syncthreads(); - - // ============================================================================ - // Passes 1-4: Refine using 8-bit passes over FP32 bits - // ============================================================================ - // FP32 bits [31:0] split into 4 bytes processed MSB-first: - // Pass 1: bits [31:24], Pass 2: bits [23:16], Pass 3: bits [15:8], Pass 4: - // bits [7:0] -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { - __shared__ int shared_final_k; // For final pass: remaining slots to fill - const int src_buffer = pass % 2; - const int dst_buffer = src_buffer ^ 1; - - // Clamp buffered count to prevent overflow - const int raw_buffered = shared_buffered_count[src_buffer]; - const int num_buffered = - (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS; - - compute_cumulative_sum(); - - // Find threshold bin for this pass - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[dst_buffer] = 0; - shared_final_k = remaining_k - shared_histogram[0][thread_id + 1]; + uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group, + static_cast(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( + logits.data_ptr(), output.data_ptr(), + lengths.data_ptr(), static_cast(num_rows), + static_cast(TopK), static_cast(stride), + stream); + TORCH_CHECK(status == cudaSuccess, + "FilteredTopK fallback failed: ", cudaGetErrorString(status)); + return; } - __syncthreads(); - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - - // Bit offset for this pass: 24, 16, 8, 0 - const int bit_offset = 24 - pass * 8; - - // Early exit if threshold bin perfectly matches - if (remaining_k == 0) { - for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { - const int idx = buffered_indices[src_buffer][i]; - const uint32_t fp32_bits = - convert_to_uint32_v2(logits[idx + logits_offset]); - const int bin = (fp32_bits >> bit_offset) & 0xFF; - if (bin > threshold_bin) { - const int output_pos = ::atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } - } - __syncthreads(); - break; + size_t state_bytes = num_groups * sizeof(P::RadixRowState); + TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), + "workspace too small, need ", state_bytes, " bytes"); + + // Zero the per-group RadixRowState region before launch — only when the + // radix path will actually run (max_seq_len > RADIX_THRESHOLD). The + // RadixRowState fields (arrival_counter, histograms) are only touched by + // radix_topk; the decode/medium paths inside the persistent kernel + // operate purely in shared memory and never read these globals, so a + // stale workspace is harmless for them. + // + // Why we need the memset (when needs_cooperative is true): + // 1. arrival_counter accumulates within a launch and is never reset, + // so a prior call leaves it at a large positive value. Without this + // reset, the very first wait_ge in the next call sees counter >> + // target and returns instantly, breaking the barrier. + // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA + // __syncthreads(), so it had no happens-before edge to CTA-1+'s + // first red_release. cudaMemsetAsync is stream-ordered: the zero + // is globally visible before any CTA runs. + if (needs_cooperative) { + cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr(), 0, + state_bytes, stream); + TORCH_CHECK(mz_err == cudaSuccess, + "row_states memset failed: ", cudaGetErrorString(mz_err)); } - // Continue refinement - __syncthreads(); - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; + P::PersistentTopKParams params; + params.input = logits.data_ptr(); + params.output = output.data_ptr(); + params.lengths = lengths.data_ptr(); + params.num_rows = static_cast(num_rows); + params.stride = static_cast(stride); + params.top_k = static_cast(TopK); + params.chunk_size = chunk_size; + params.row_states = + reinterpret_cast(workspace.data_ptr()); + params.ctas_per_group = ctas_per_group; + params.max_seq_len = static_cast(max_seq_len); + + #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ + do { \ + auto kernel = &P::persistent_topk_kernel; \ + cudaError_t err = cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ + TORCH_CHECK(err == cudaSuccess, \ + "Failed to set smem: ", cudaGetErrorString(err)); \ + kernel<<>>(params); \ + } while (0) + + if (vec_size == 4) { + LAUNCH_PERSISTENT(TopK, 4); + } else if (vec_size == 2) { + LAUNCH_PERSISTENT(TopK, 2); + } else { + LAUNCH_PERSISTENT(TopK, 1); } - __syncthreads(); - - for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { - const int idx = buffered_indices[src_buffer][i]; - const float logit_value = logits[idx + logits_offset]; - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int bin = (fp32_bits >> bit_offset) & 0xFF; - - if (bin > threshold_bin) { - // Definitely in top-k - const int output_pos = ::atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - if (pass == 3) { - // Final pass (bits [7:0]): No more refinement possible - // Fill remaining slots in reverse order to maintain descending order - const int slot = ::atomicAdd(&shared_final_k, -1); - if (slot > 0) { - output_indices[TopK - slot] = idx; - } - } else { - // Buffer for next pass and build next histogram - const int buffer_pos = - ::atomicAdd(&shared_buffered_count[dst_buffer], 1); - if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) { - buffered_indices[dst_buffer][buffer_pos] = idx; - // Fused: Build histogram for next pass - const int next_bit_offset = bit_offset - 8; - const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF; - ::atomicAdd(&shared_histogram[0][next_bin], 1); - } - } - } - } - __syncthreads(); + #undef LAUNCH_PERSISTENT } -} -__global__ __launch_bounds__(kThreadsPerBlock) void topk_kernel( - const FastTopKParams params) { - const auto& [input, row_starts, indices, lengths, input_stride] = params; - const uint64_t batch_idx = blockIdx.x; - const int logits_offset = row_starts == nullptr ? 0 : row_starts[batch_idx]; - const int seq_len = lengths[batch_idx]; - int* output_indices = indices + batch_idx * TopK; - const float* logits = input + batch_idx * input_stride; - - if (seq_len <= TopK) { - // Shortcut: All elements are in top-k - return naive_topk_cuda(logits, output_indices, seq_len); - } else { - return fast_topk_cuda_tl(logits, output_indices, logits_offset, seq_len); - } + cudaError_t err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, + "persistent_topk failed: ", cudaGetErrorString(err)); } +#endif -FastTopKParams get_params( - const at::Tensor& score, const at::Tensor& lengths, - std::optional row_starts_opt = std::nullopt, - std::optional indices_opt = std::nullopt) { - const int64_t batch_size = score.size(0); - - TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1, - "score must be 2D with contiguous rows"); - TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous() && - lengths.size(0) == batch_size, - "lengths must be 1D contiguous with size matching batch"); - - const int32_t* row_starts_ptr = nullptr; - if (row_starts_opt.has_value()) { - const auto& row_starts = *row_starts_opt; - TORCH_CHECK(row_starts.dim() == 1 && row_starts.size(0) == batch_size, - "row_starts must be 1D with size matching batch"); - row_starts_ptr = row_starts.data_ptr(); - } +} // anonymous namespace - int32_t* indices_ptr = nullptr; - if (indices_opt.has_value()) { - const auto& indices = *indices_opt; - TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous() && - indices.size(0) == batch_size && indices.size(1) == TopK, - "indices must be 2D contiguous [batch, TopK]"); - indices_ptr = indices.data_ptr(); +void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, + torch::Tensor& output, torch::Tensor& workspace, int64_t k, + int64_t max_seq_len) { +#ifndef USE_ROCM + TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); + TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); + TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); + TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported"); + TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32"); + TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, + "lengths must be 1D or 2D"); + TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); + TORCH_CHECK(output.dim() == 2, "output must be 2D"); + + const int64_t num_rows = logits.size(0); + const int64_t stride = logits.size(1); + + TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); + TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, + "output size mismatch"); + TORCH_CHECK(k == 512 || k == 1024 || k == 2048, + "persistent_topk supports k=512, k=1024, or k=2048, got k=", k); + + if (k == 512) { + launch_persistent_topk<512>(logits, lengths, output, workspace, + max_seq_len); + } else if (k == 1024) { + launch_persistent_topk<1024>(logits, lengths, output, workspace, + max_seq_len); + } else { + launch_persistent_topk<2048>(logits, lengths, output, workspace, + max_seq_len); } - - return FastTopKParams{ - .input = score.data_ptr(), - .row_starts = row_starts_ptr, - .indices = indices_ptr, - .lengths = lengths.data_ptr(), - .input_stride = score.stride(0), - }; -} - -template -void setup_kernel_smem_once() { - static const cudaError_t result = []() -> cudaError_t { -#ifdef USE_ROCM - auto func_ptr = reinterpret_cast(kernel_func); #else - auto func_ptr = kernel_func; + TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); #endif - return cudaFuncSetAttribute( - func_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - }(); - - TORCH_CHECK( - result == cudaSuccess, - "Failed to set kernel shared memory limit: ", cudaGetErrorString(result)); } - -} // namespace vllm - -void large_context_topk( - const torch::Tensor& logits, torch::Tensor& indices, - const torch::Tensor& seq_lens, - std::optional row_starts = std::nullopt) { - TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); - TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor"); - TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor"); - if (row_starts.has_value()) { - TORCH_CHECK(row_starts->is_cuda(), "row_starts must be a CUDA tensor"); - } - - const auto params = vllm::get_params(logits, seq_lens, row_starts, indices); - const int64_t batch_size = logits.size(0); - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const dim3 grid(static_cast(batch_size)); - const dim3 block(vllm::kThreadsPerBlock); - - vllm::setup_kernel_smem_once(); - vllm::topk_kernel<<>>(params); - - const cudaError_t result = cudaGetLastError(); - TORCH_CHECK(result == cudaSuccess, - "large_context_topk kernel failed: ", cudaGetErrorString(result)); -} \ No newline at end of file diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3593f1d2225b..8d8f7bed0441 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -106,10 +106,28 @@ 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); + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + ops.impl("silu_and_mul_per_block_quant", torch::kCUDA, + &silu_and_mul_per_block_quant); + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu); @@ -161,9 +179,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " "int num_heads_k, int num_heads_v, int head_dim, float eps, " "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids) -> ()"); + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope); +#ifndef USE_ROCM + // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and + // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one + // kernel launch. + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor! q, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "float eps, int cache_block_size) -> ()"); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA, + &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert); +#endif + // Apply repetition penalties to logits in-place ops.def( "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " @@ -185,10 +217,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode); ops.def( - "large_context_topk(Tensor score, Tensor indices, Tensor lengths, " - "Tensor? " - "row_starts_opt) -> ()"); - ops.impl("large_context_topk", torch::kCUDA, &large_context_topk); + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + ops.impl("persistent_topk", torch::kCUDA, &persistent_topk); // Layernorm-quant // Apply Root Mean Square (RMS) Normalization to the input tensor. @@ -228,22 +259,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "rotary_embedding(Tensor positions, Tensor! query," " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox) -> ()"); + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding); // Quantization ops #ifndef USE_ROCM - // Fused SiLU+Mul + per-block quantization - ops.def( - "silu_and_mul_per_block_quant(" - "Tensor! out, " - "Tensor input, " - "Tensor! scales, " - "int group_size, " - "Tensor? scale_ub=None, " - "bool is_scale_transposed=False) -> ()"); - ops.impl("silu_and_mul_per_block_quant", torch::kCUDA, - &silu_and_mul_per_block_quant); // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). ops.def( "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); @@ -496,6 +517,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? b_qzeros, " "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt " "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor"); + + ops.def( + "minimax_allreduce_rms(" + "Tensor input," + "Tensor norm_weight," + "Tensor workspace," + "int rank," + "int nranks," + "float eps) -> Tensor"); + ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms); + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv," + "Tensor norm_weight_q," + "Tensor norm_weight_k," + "Tensor workspace," + "int q_size," + "int kv_size," + "int rank," + "int nranks," + "float eps) -> (Tensor, Tensor)"); + ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk); + // conditionally compiled so impl in source file #endif } diff --git a/docker/Dockerfile b/docker/Dockerfile index 6bbd34f95437..fc2c919eea68 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -22,7 +22,7 @@ # docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json # ============================================================================= -ARG CUDA_VERSION=13.0.0 +ARG CUDA_VERSION=13.0.2 ARG PYTHON_VERSION=3.12 ARG UBUNTU_VERSION=22.04 @@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Explicitly set the list to avoid issues with torch 2.2 # See https://github.com/pytorch/pytorch/pull/123243 # From versions.json: .torch.cuda_arch_list -ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0' +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} #################### BUILD BASE IMAGE #################### @@ -204,7 +204,7 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL ARG PYTORCH_NIGHTLY # Install build dependencies -COPY requirements/build.txt requirements/build.txt +COPY requirements/build/cuda.txt requirements/build/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -219,13 +219,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing build requirements without torch..." \ && python3 use_existing_torch.py --prefix \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ && echo "Installing torch nightly..." \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ echo "Installing build requirements..." \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi @@ -315,7 +315,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ #################### CSRC BUILD IMAGE #################### #################### EXTENSIONS BUILD IMAGE #################### -# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build +# Build DeepEP - runs in PARALLEL with csrc-build # This stage is independent and doesn't affect csrc cache FROM base AS extensions-build ARG CUDA_VERSION @@ -327,21 +327,6 @@ ENV UV_LINK_MODE=copy WORKDIR /workspace -# Build DeepGEMM wheel -# Default moved here from tools/install_deepgemm.sh for centralized version management -ARG DEEPGEMM_GIT_REF=477618cd51baffca09c4b0b87e97c03fe827ef03 -COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh -RUN --mount=type=cache,target=/root/.cache/uv \ - mkdir -p /tmp/deepgemm/dist && \ - VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \ - --cuda-version "${CUDA_VERSION}" \ - ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \ - --wheel-dir /tmp/deepgemm/dist || \ - echo "DeepGEMM build skipped (CUDA version requirement not met)" - -# Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped -RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped - # Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management @@ -370,7 +355,7 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL ARG PYTORCH_NIGHTLY # Install build dependencies -COPY requirements/build.txt requirements/build.txt +COPY requirements/build/cuda.txt requirements/build/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -385,13 +370,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing build requirements without torch..." \ && python3 use_existing_torch.py --prefix \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ && echo "Installing torch nightly..." \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ echo "Installing build requirements..." \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi @@ -426,7 +411,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 # Copy extension wheels from extensions-build stage for later use -COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist # Check the size of the wheel if RUN_WHEEL_CHECK is true @@ -466,8 +450,8 @@ ARG PYTORCH_NIGHTLY # Install development dependencies COPY requirements/lint.txt requirements/lint.txt -COPY requirements/test.in requirements/test.in -COPY requirements/test.txt requirements/test.txt +COPY requirements/test/cuda.in requirements/test/cuda.in +COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -475,8 +459,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing dev requirements plus torch nightly..." \ && python3 use_existing_torch.py --prefix \ - && cat torch_lib_versions.txt >> requirements/test.in \ - && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \ + && cat torch_lib_versions.txt >> requirements/test/cuda.in \ + && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | xargs) --pre \ -r requirements/dev.txt \ @@ -506,41 +490,46 @@ WORKDIR /vllm-workspace RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \ echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment -# Install Python and system dependencies +# Install Python from source and system dependencies RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ - software-properties-common \ curl \ sudo \ ffmpeg \ libsm6 \ libxext6 \ libgl1 \ - && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \ - if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \ - mkdir -p -m 0755 /etc/apt/keyrings ; \ - curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \ - sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \ - echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \ - fi ; \ - else \ - for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ - done ; \ - fi \ - && apt-get update -y \ - && apt-get install -y --no-install-recommends \ - python${PYTHON_VERSION} \ - python${PYTHON_VERSION}-dev \ - python${PYTHON_VERSION}-venv \ libibverbs-dev \ + build-essential \ + libssl-dev \ + libffi-dev \ + zlib1g-dev \ + libbz2-dev \ + libreadline-dev \ + libsqlite3-dev \ + libncurses-dev \ + liblzma-dev \ + libgdbm-dev \ + uuid-dev \ + tk-dev \ + && PYTHON_MAJOR_MINOR=${PYTHON_VERSION} \ + && PYTHON_FULL_VERSION=$(curl -s https://www.python.org/ftp/python/ \ + | grep -oE "${PYTHON_MAJOR_MINOR}\.[0-9]+" \ + | sort -t. -k3 -n | uniq \ + | tail -1) \ + && echo "Building Python ${PYTHON_FULL_VERSION} from source..." \ + && curl -fSL https://www.python.org/ftp/python/${PYTHON_FULL_VERSION}/Python-${PYTHON_FULL_VERSION}.tgz -o /tmp/python.tgz \ + && tar -xzf /tmp/python.tgz -C /tmp \ + && cd /tmp/Python-${PYTHON_FULL_VERSION} \ + && ./configure --enable-optimizations --with-ensurepip=install --prefix=/usr/local \ + && make -j$(nproc) \ + && make install \ + && cd / \ + && rm -rf /tmp/python.tgz /tmp/Python-${PYTHON_FULL_VERSION} \ + && ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR} /usr/bin/python3 \ + && ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR}-config /usr/bin/python3-config \ + && ln -sf /usr/local/bin/pip${PYTHON_MAJOR_MINOR} /usr/bin/pip \ && rm -rf /var/lib/apt/lists/* \ - && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ - && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ - && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \ - && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \ && python3 --version && python3 -m pip --version # Install CUDA development tools for runtime JIT compilation @@ -554,7 +543,11 @@ 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-${CUDA_VERSION_DASH} && \ + libcublas-${CUDA_VERSION_DASH} \ + # Required by fastsafetensors (fixes #20384) + libnuma-dev \ + # numactl CLI for NUMA binding at runtime + numactl && \ # 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, @@ -594,33 +587,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798) -# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on -# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738 -ARG FLASHINFER_VERSION=0.6.7 +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 '.') \ - && flashinfer show-config - -# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments. -# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com -# which fails without internet. This step caches them at build time. -RUN python3 <<'PYEOF' -from flashinfer.jit import env as jit_env -from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin -from flashinfer.artifacts import ArtifactPath, CheckSumHash - -download_trtllm_headers( - 'bmm', - jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export', - f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export', - ArtifactPath.TRTLLM_GEN_BMM, - get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM), -) - -print('FlashInfer TRTLLM BMM headers downloaded successfully') -PYEOF + && flashinfer show-config \ + && flashinfer download-cubin # ============================================================ # OPENAI API SERVER DEPENDENCIES @@ -656,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ else \ BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \ fi; \ - uv pip install --system accelerate hf_transfer modelscope \ + uv pip install --system accelerate modelscope \ "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}" # ============================================================ @@ -693,15 +665,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ . /etc/environment && \ uv pip list -# Install deepgemm wheel that has been built in the `build` stage -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \ - sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \ - uv pip install --system /tmp/deepgemm/dist/*.whl; \ - else \ - echo "No DeepGEMM wheels to install; skipping."; \ - fi' - # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH @@ -750,8 +713,8 @@ ARG PYTORCH_NIGHTLY # Install development dependencies (for testing) COPY requirements/lint.txt requirements/lint.txt -COPY requirements/test.in requirements/test.in -COPY requirements/test.txt requirements/test.txt +COPY requirements/test/cuda.in requirements/test/cuda.in +COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -761,8 +724,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing dev requirements plus torch nightly..." \ && python3 use_existing_torch.py --prefix \ - && cat torch_lib_versions.txt >> requirements/test.in \ - && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \ + && cat torch_lib_versions.txt >> requirements/test/cuda.in \ + && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \ -r requirements/dev.txt \ @@ -779,9 +742,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -e tests/vllm_test_utils # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER 1 +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 # Copy in the v1 package for testing (it isn't distributed yet) COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 @@ -808,7 +772,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL ENV UV_HTTP_TIMEOUT=500 # install kv_connectors if requested -ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0' +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index afcb388beb29..72ee002d90a4 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -107,10 +107,10 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \ fi # Copy build requirements -COPY requirements/cpu-build.txt requirements/build.txt +COPY requirements/build/cpu.txt requirements/build/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/build.txt + uv pip install -r requirements/build/cpu.txt COPY . . @@ -127,26 +127,28 @@ FROM base AS vllm-test-deps WORKDIR /vllm-workspace # Copy test requirements -COPY requirements/test.in requirements/cpu-test.in +COPY requirements/test/cuda.in requirements/test/cpu.in RUN \ - sed -i '/mamba_ssm/d' requirements/cpu-test.in && \ + sed -i '/mamba_ssm/d' requirements/test/cpu.in && \ remove_packages_not_supported_on_aarch64() { \ case "$(uname -m)" in \ aarch64|arm64) \ - sed -i '/decord/d' requirements/cpu-test.in; \ - sed -i '/terratorch/d' requirements/cpu-test.in; \ + sed -i '/decord/d' requirements/test/cpu.in; \ + sed -i '/terratorch/d' requirements/test/cpu.in; \ ;; \ esac; \ }; \ remove_packages_not_supported_on_aarch64 && \ - sed -i 's/^torch==.*/torch==2.11.0/g' requirements/cpu-test.in && \ - sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \ - sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \ - uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu + sed -i 's/^torch==.*/torch==2.11.0/g' requirements/test/cpu.in && \ + sed -i 's/torchaudio.*/torchaudio/g' requirements/test/cpu.in && \ + sed -i 's/torchvision.*/torchvision/g' requirements/test/cpu.in && \ + # Related issue: https://github.com/vllm-project/vllm/pull/38800#issuecomment-4228314305 + sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/cpu.in && \ + uv pip compile requirements/test/cpu.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/cpu-test.txt + uv pip install -r requirements/test/cpu.txt ######################### DEV IMAGE ######################### FROM vllm-build AS vllm-dev @@ -168,10 +170,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=.git,target=.git \ VLLM_TARGET_DEVICE=cpu python3 setup.py develop -COPY --from=vllm-test-deps /vllm-workspace/requirements/cpu-test.txt requirements/test.txt +COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/dev.txt && \ + uv pip install -r requirements/lint.txt && \ + uv pip install -r requirements/test/cpu.txt && \ pre-commit install --hook-type pre-commit --hook-type commit-msg ENTRYPOINT ["bash"] @@ -195,6 +198,12 @@ ADD ./.buildkite/ ./.buildkite/ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -e tests/vllm_test_utils +# enable fast downloads from hf (for testing) +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 + ######################### RELEASE IMAGE ######################### FROM base AS vllm-openai diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 045a09a42dc4..6b9d85c9d17f 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128 RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system numba==0.61.2 + uv pip install --system numba==0.65.0 RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -r requirements/common.txt @@ -107,7 +107,7 @@ COPY . . RUN python3 use_existing_torch.py RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/build.txt + uv pip install --system -r requirements/build/cuda.txt ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ @@ -217,16 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2. # build flashinfer for torch nightly from source around 10 mins -# release version: v0.6.7 -# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798) -# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on -# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738 +# release version: v0.6.8.post1 # todo(elainewy): cache flashinfer build result for faster build ENV CCACHE_DIR=/root/.cache/ccache RUN --mount=type=cache,target=/root/.cache/ccache \ --mount=type=cache,target=/root/.cache/uv \ echo "git clone flashinfer..." \ - && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \ + && git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \ && cd flashinfer \ && git submodule update --init --recursive \ && echo "finish git clone flashinfer..." \ @@ -261,7 +258,7 @@ FROM vllm-base as test COPY tests/ tests/ # install build and runtime dependencies without stable torch version -COPY requirements/nightly_torch_test.txt requirements/nightly_torch_test.txt +COPY requirements/test/nightly-torch.txt requirements/test/nightly-torch.txt # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out # Reference: https://github.com/astral-sh/uv/pull/1694 @@ -272,12 +269,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system -e tests/vllm_test_utils # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER 1 +ENV HF_XET_HIGH_PERFORMANCE 1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/nightly_torch_test.txt + uv pip install --system -r requirements/test/nightly-torch.txt # Logging to confirm the torch versions RUN pip freeze | grep -E 'torch|vllm|flashinfer' diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 07b64a509a4b..845d900c39cd 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -251,7 +251,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ make -C /numactl install && \ # sentencepiece.pc is in some pkgconfig inside uv cache export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':') && \ - nanobind_DIR=$(uv pip show nanobind | grep Location | sed 's/^Location: //;s/$/\/nanobind\/cmake/') && uv pip install -r /src/requirements/common.txt -r /src/requirements/cpu.txt -r /src/requirements/build.txt --no-build-isolation && \ + nanobind_DIR=$(uv pip show nanobind | grep Location | sed 's/^Location: //;s/$/\/nanobind\/cmake/') && uv pip install -r /src/requirements/common.txt -r /src/requirements/cpu.txt -r /src/requirements/build/cuda.txt --no-build-isolation && \ cd /src/ && \ uv build --wheel --out-dir /vllmwheel/ --no-build-isolation && \ uv pip install /vllmwheel/*.whl diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fad79af9cffd..7a93823cd2d5 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -2,6 +2,11 @@ ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base +# AMD NIC backend +ARG NIC_BACKEND=none +# AMD AINIC apt repo settings +ARG AINIC_VERSION=1.117.5 +ARG UBUNTU_CODENAME=jammy # Sccache configuration (only used in release pipeline) ARG USE_SCCACHE @@ -19,7 +24,8 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install some basic utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ - apt-transport-https ca-certificates wget curl + apt-transport-https ca-certificates wget curl \ + libnuma-dev RUN python3 -m pip install --upgrade pip # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE @@ -118,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages FROM base AS build_rixl -ARG RIXL_BRANCH="f33a5599" +ARG RIXL_BRANCH="bf4a7214" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" -ARG UCX_BRANCH="da3fac2a" -ARG UCX_REPO="https://github.com/ROCm/ucx.git" +ARG UCX_BRANCH="7009d7a1" +ARG UCX_REPO="https://github.com/openucx/ucx.git" ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl @@ -159,7 +165,7 @@ RUN cd /usr/local/src && \ --disable-doxygen-doc \ --enable-optimizations \ --enable-devel-headers \ - --with-rocm=/opt/rocm \ + --with-rocm=${ROCM_PATH} \ --with-verbs \ --with-dm \ --enable-mt && \ @@ -180,7 +186,12 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \ ninja install # Generate RIXL wheel -RUN cd /opt/rixl && mkdir -p /app/install && \ +# Exclude libcore and libpull from auditwheel: transitive dependencies +# that are not shipped in the wheel and vary across base images. +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + mkdir -p /app/install && \ ./contrib/build-wheel.sh \ --output-dir /app/install \ --rocm-dir ${ROCM_PATH} \ @@ -191,9 +202,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \ FROM base AS build_deep ARG ROCSHMEM_BRANCH="ba0bf0f3" ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -ARG DEEPEP_BRANCH="e84464ec" +ARG DEEPEP_BRANCH="5d90af8b" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" ENV ROCSHMEM_DIR=/opt/rocshmem RUN git clone ${ROCSHMEM_REPO} \ @@ -201,13 +213,11 @@ RUN git clone ${ROCSHMEM_REPO} \ && git checkout ${ROCSHMEM_BRANCH} \ && mkdir -p projects/rocshmem/build \ && cd projects/rocshmem/build \ - && cmake .. \ - -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \ - -DROCM_PATH=/opt/rocm \ - -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ - -DUSE_EXTERNAL_MPI=OFF \ - && make -j \ - && make install + && 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. @@ -216,6 +226,47 @@ RUN git clone ${DEEPEP_REPO} \ && git checkout ${DEEPEP_BRANCH} \ && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install +# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do +# not force users to rebuild the long-lived Dockerfile.rocm_base image. +FROM base AS mori_base +ARG NIC_BACKEND +ARG AINIC_VERSION +ARG UBUNTU_CODENAME +RUN /bin/bash -lc 'set -euo pipefail; \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). + # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed. + case "${NIC_BACKEND}" in \ + # default: mlx5 + none) \ + ;; \ + # AMD NIC + ainic) \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \ + rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + ;; \ + # TODO: Add Broadcom bnxt packages/repos here later. + # bnxt) \ + # echo "[MORI] Add Broadcom bnxt packages/repos here later."; \ + # ;; \ + *) \ + echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \ + exit 2; \ + ;; \ + esac;' + # ----------------------- # vLLM wheel release build stage (for building distributable wheels) # This stage pins dependencies to custom ROCm wheel versions and handles version detection @@ -318,7 +369,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- # Test vLLM image -FROM base AS test +FROM mori_base AS test RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* @@ -328,14 +379,14 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ && uv pip install --system -r requirements/rocm.txt \ - && uv pip install --system -r requirements/rocm-test.txt \ + && uv pip install --system -r requirements/test/rocm.txt \ && pip uninstall -y vllm \ && uv pip install --system *.whl -# Verify that PyTorch is the ROCm build, not CUDA -RUN python3 -c "import torch; assert torch.version.hip is not None, \ - f'Expected ROCm PyTorch but got CUDA (torch.version.cuda={torch.version.cuda}, torch.version.hip={torch.version.hip})'; \ - print(f'Verified: PyTorch {torch.__version__} with ROCm (HIP {torch.version.hip})')" +# Persist the built wheel in the image so python_only_compile_rocm.sh can +# reinstall it after removing compilers. The bind-mounted /install contents +# above are not available once that RUN step completes. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ # Install RIXL wheel RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ @@ -364,9 +415,10 @@ RUN cd /vllm-workspace \ && python3 -m pip install pytest-shard # enable fast downloads from hf (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system hf_transfer -ENV HF_HUB_ENABLE_HF_TRANSFER=1 +ENV HF_XET_HIGH_PERFORMANCE=1 + +# increase timeout for hf downloads (for testing) +ENV HF_HUB_DOWNLOAD_TIMEOUT 60 # install audio decode package `torchcodec` from source (required due to # ROCm and torch version mismatch) for tests with datasets package @@ -384,6 +436,10 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac ENV MIOPEN_DEBUG_CONV_DIRECT=0 ENV MIOPEN_DEBUG_CONV_GEMM=0 +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + # Source code is used in the `python_only_compile.sh` test # We hide it inside `src/` so that this source code # will not be imported by other tests @@ -408,7 +464,7 @@ RUN printf '%s\n' \ # ----------------------- # Final vLLM image -FROM base AS final +FROM mori_base AS final RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* @@ -445,6 +501,8 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ ARG COMMON_WORKDIR ARG BASE_IMAGE +ARG NIC_BACKEND +ARG AINIC_VERSION # Copy over the benchmark scripts as well COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks @@ -462,7 +520,9 @@ ENV HIP_FORCE_DEV_KERNARG=1 # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt CMD ["/bin/bash"] diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index e77406728cb4..5940a4ee564d 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,9 +9,9 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.10.post2" +ARG AITER_BRANCH="v0.1.10.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" -ARG MORI_BRANCH="2d02c6a9" +ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" # Sccache configuration (only used in release pipeline) @@ -112,10 +112,14 @@ FROM base AS build_triton ARG TRITON_BRANCH ARG TRITON_REPO RUN git clone ${TRITON_REPO} +# Cherry picking the following +# https://github.com/triton-lang/triton/pull/8991 +# https://github.com/triton-lang/triton/pull/9541 RUN cd triton \ && git checkout ${TRITON_BRANCH} \ && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ && git cherry-pick 555d04f \ + && git cherry-pick dd998b6 \ && if [ ! -f setup.py ]; then cd python; fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ && mkdir -p /app/install && cp dist/*.whl /app/install diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index e90f2fdfc4c3..554a7257c236 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -42,7 +42,7 @@ FROM python-install AS pyarrow # Build Apache Arrow WORKDIR /tmp RUN --mount=type=cache,target=/root/.cache/uv \ - git clone https://github.com/apache/arrow.git && \ + git clone https://github.com/apache/arrow.git -b maint-19.0.1 && \ cd arrow/cpp && \ mkdir release && cd release && \ cmake -DCMAKE_BUILD_TYPE=Release \ @@ -68,19 +68,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements-build.txt && \ python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel -FROM python-install AS numa-build -# Install numactl (needed for numa.h dependency) -WORKDIR /tmp -RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \ - tar -xvzf v2.0.16.tar.gz && \ - cd numactl-2.0.16 && \ - ./autogen.sh && \ - ./configure && \ - make - -# Set include path -ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" - FROM python-install AS rust ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup @@ -91,6 +78,18 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \ rustup default stable && \ rustup show +FROM python-install AS numa-build +WORKDIR /tmp +RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \ + tar -xvzf v2.0.19.tar.gz && \ + cd numactl-2.0.19 && \ + ./autogen.sh && \ + ./configure && \ + make + +# Set include path +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" + FROM python-install AS torch-vision # Install torchvision ARG TORCH_VISION_VERSION=v0.26.0 @@ -133,7 +132,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \ git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ cd llvm-project && mkdir build && cd build && \ - uv pip install 'cmake<4' setuptools numpy && \ + uv pip install 'cmake<4' 'setuptools<70' numpy && \ export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \ CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \ CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \ @@ -193,27 +192,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \ cd opencv-python && \ python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist -# Build Outlines Core -FROM python-install AS outlines-core-builder +## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330 +FROM python-install AS guidance-builder WORKDIR /tmp ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH" -COPY requirements/common.txt /tmp/requirements/common.txt -ARG OUTLINES_CORE_VERSION RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \ --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \ - OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \ - if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \ - git clone https://github.com/dottxt-ai/outlines-core.git && \ - cd outlines-core && \ - git checkout tags/${OUTLINES_CORE_VERSION} && \ - sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \ + git clone https://github.com/guidance-ai/llguidance.git && \ + cd llguidance && \ + git checkout s390x-fix-v2 && \ uv pip install maturin && \ - python -m maturin build --release --out dist + python -m maturin build --release --out dist --compatibility linux -# Final build stage +# # Final build stage FROM python-install AS vllm-cpu ARG PYTHON_VERSION ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" @@ -229,10 +223,12 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/ ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH" ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +# Force pure Python protobuf to avoid s390x C++ extension crashes +ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python COPY . /workspace/vllm WORKDIR /workspace/vllm -RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \ +RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \ make -C /numactl install # Install dependencies, including PyTorch and Apache Arrow @@ -245,24 +241,24 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \ --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \ --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \ - --mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \ - ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \ + --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \ + ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \ VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \ HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \ LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \ NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \ OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \ - OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \ - uv pip install -v \ - $ARROW_WHL_FILE \ + GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \ + uv pip install -v \ + $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ $LLVM_WHL_FILE \ $NUMBA_WHL_FILE \ $OPENCV_WHL_FILE \ - $OUTLINES_CORE_WHL_FILE \ + $GUIDANCE_WHL_FILE \ --index-strategy unsafe-best-match \ - -r requirements/cpu-build.txt \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt @@ -271,6 +267,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \ uv pip install "$(echo dist/*.whl)[tensorizer]" +# Remove protobuf C++ extension that crashes on s390x +RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \ + /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true + # setup non-root user for vllm RUN umask 002 && \ /usr/sbin/useradd --uid 2000 --gid 0 vllm && \ diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index b8ec60a82324..ab3f6f40d87b 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} ENV PATH="$VIRTUAL_ENV/bin:$PATH" -# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2. -ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh" -RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \ +# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3. +ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" +RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \ bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \ rm "${ONECCL_INSTALLER}" && \ echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ @@ -76,20 +76,14 @@ ENV UV_LINK_MODE="copy" RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \ --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \ - --mount=type=bind,src=requirements/xpu-test.in,target=/workspace/vllm/requirements/xpu-test.in \ + --mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \ uv pip install --upgrade pip && \ uv pip install -r requirements/xpu.txt && \ - uv pip compile /workspace/vllm/requirements/xpu-test.in \ - -o /workspace/vllm/requirements/xpu-test.txt \ - -c /workspace/vllm/requirements/xpu.txt \ - --index-strategy unsafe-best-match \ - --extra-index-url ${PIP_EXTRA_INDEX_URL} \ - --python-version ${PYTHON_VERSION} && \ uv pip install grpcio-tools protobuf nanobind && \ source /opt/intel/oneapi/setvars.sh --force && \ source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \ - uv pip install --no-build-isolation -r /workspace/vllm/requirements/xpu-test.txt + uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt @@ -170,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # FIX triton RUN --mount=type=cache,target=/root/.cache/uv \ uv pip uninstall triton triton-xpu && \ - uv pip install triton-xpu==3.6.0 + uv pip install triton-xpu==3.7.0 # remove torch bundled oneccl to avoid conflicts RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/docker/docker-bake.hcl b/docker/docker-bake.hcl index e1c2fbba63a6..785d598d6080 100644 --- a/docker/docker-bake.hcl +++ b/docker/docker-bake.hcl @@ -20,7 +20,7 @@ variable "NVCC_THREADS" { } variable "TORCH_CUDA_ARCH_LIST" { - default = "8.0 8.9 9.0 10.0" + default = "8.0 8.9 9.0 10.0 11.0 12.0" } variable "COMMIT" { @@ -88,7 +88,6 @@ target "test-ubuntu2404" { args = { UBUNTU_VERSION = "24.04" GDRCOPY_OS_VERSION = "Ubuntu24_04" - FLASHINFER_AOT_COMPILE = "true" } output = ["type=docker"] } @@ -100,7 +99,6 @@ target "openai-ubuntu2404" { args = { UBUNTU_VERSION = "24.04" GDRCOPY_OS_VERSION = "Ubuntu24_04" - FLASHINFER_AOT_COMPILE = "true" } output = ["type=docker"] } diff --git a/docker/versions.json b/docker/versions.json index 71caab20b0c8..b6b555790d2a 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -2,7 +2,7 @@ "_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py", "variable": { "CUDA_VERSION": { - "default": "13.0.0" + "default": "13.0.2" }, "PYTHON_VERSION": { "default": "3.12" @@ -11,10 +11,10 @@ "default": "22.04" }, "BUILD_BASE_IMAGE": { - "default": "nvidia/cuda:13.0.0-devel-ubuntu22.04" + "default": "nvidia/cuda:13.0.2-devel-ubuntu22.04" }, "FINAL_BASE_IMAGE": { - "default": "nvidia/cuda:13.0.0-base-ubuntu22.04" + "default": "nvidia/cuda:13.0.2-base-ubuntu22.04" }, "GET_PIP_URL": { "default": "https://bootstrap.pypa.io/get-pip.py" @@ -32,7 +32,7 @@ "default": "false" }, "TORCH_CUDA_ARCH_LIST": { - "default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0" + "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX" }, "MAX_JOBS": { "default": "2" @@ -52,9 +52,6 @@ "vllm_target_device": { "default": "cuda" }, - "DEEPGEMM_GIT_REF": { - "default": "477618cd51baffca09c4b0b87e97c03fe827ef03" - }, "DEEPEP_COMMIT_HASH": { "default": "73b6ea4" }, @@ -68,7 +65,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.7" + "default": "0.6.8.post1" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/README.md b/docs/README.md index 4b480c463abb..2bdc3e1a3e4c 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,7 +25,7 @@ hide: vLLM is a fast and easy-to-use library for LLM inference and serving. -Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has evolved into a community-driven project with contributions from both academia and industry. +Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors. Where to get started with vLLM depends on the type of user. If you are looking to: @@ -42,23 +42,37 @@ vLLM is fast with: - State-of-the-art serving throughput - Efficient management of attention key and value memory with [**PagedAttention**](https://blog.vllm.ai/2023/06/20/vllm.html) -- Continuous batching of incoming requests -- Fast model execution with CUDA/HIP graph -- Quantization: [GPTQ](https://arxiv.org/abs/2210.17323), [AWQ](https://arxiv.org/abs/2306.00978), INT4, INT8, and FP8 -- Optimized CUDA kernels, including integration with FlashAttention and FlashInfer. -- Speculative decoding -- Chunked prefill +- Continuous batching of incoming requests, chunked prefill, prefix caching +- Fast and flexible model execution with piecewise and full CUDA/HIP graphs +- Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and [more](https://docs.vllm.ai/en/latest/features/quantization/index.html) +- Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton +- Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL +- Speculative decoding including n-gram, suffix, EAGLE, DFlash +- Automatic kernel generation and graph-level transformations using torch.compile +- Disaggregated prefill, decode, and encode vLLM is flexible and easy to use with: -- Seamless integration with popular HuggingFace models +- Seamless integration with popular Hugging Face models - High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more -- Tensor, pipeline, data and expert parallelism support for distributed inference +- Tensor, pipeline, data, expert, and context parallelism for distributed inference - Streaming outputs -- OpenAI-compatible API server -- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, Arm CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend. -- Prefix caching support -- Multi-LoRA support +- Generation of structured outputs using xgrammar or guidance +- Tool calling and reasoning parsers +- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support +- Efficient multi-LoRA support for dense and MoE layers +- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. + +vLLM seamlessly supports 200+ model architectures on HuggingFace, including: + +- Decoder-only LLMs (e.g., Llama, Qwen, Gemma) +- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS) +- Hybrid attention and state-space models (e.g., Mamba, Qwen3.5) +- Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral) +- Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT) +- Reward and classification models (e.g., Qwen-Math) + +Find the full list of supported models [here](./models/supported_models.md). For more information, check out the following: diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 5ea354f34a03..d9d501319595 100644 Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ diff --git a/docs/assets/contributing/vllm_bench_serve_dataset_stats.png b/docs/assets/contributing/vllm_bench_serve_dataset_stats.png new file mode 100644 index 000000000000..72c19d3d7c07 Binary files /dev/null and b/docs/assets/contributing/vllm_bench_serve_dataset_stats.png differ diff --git a/docs/assets/contributing/vllm_bench_serve_timeline.html b/docs/assets/contributing/vllm_bench_serve_timeline.html new file mode 100644 index 000000000000..d463e202b6d1 --- /dev/null +++ b/docs/assets/contributing/vllm_bench_serve_timeline.html @@ -0,0 +1,3888 @@ + + + +
+
+ + \ No newline at end of file diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index f78ae8a95366..c9ceb67cce67 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -37,6 +37,7 @@ th { | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` | | HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` | | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` | +| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` | | Custom | ✅ | ✅ | Local file: `data.jsonl` | | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` | @@ -107,6 +108,38 @@ P99 ITL (ms): 8.39 ================================================== ``` +#### Results Visualization + +The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis. + +```bash +vllm bench serve \ + --backend vllm \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --endpoint /v1/completions \ + --dataset-name sharegpt \ + --dataset-path /ShareGPT_V3_unfiltered_cleaned_split.json \ + --num-prompts 100 \ + --plot-timeline \ + --timeline-itl-thresholds 2,5 \ + --plot-dataset-stats \ + --save-result +``` + +##### Interactive Timeline + +The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms) + +Example output: + + + +##### Dataset statistics + +The generated figure shows the input prompt and output tokens distribution. + +Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_dataset_stats.png) + #### Custom Dataset If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl @@ -239,6 +272,69 @@ vllm bench serve \ --spec-bench-category "summarization" ``` +#### SPEED-Bench Benchmark with Speculative Decoding + +[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k). + +!!! note + This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets. + +First, download the dataset to a folder, using this one liner: + +```bash +curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 - +``` + +The command supports also the following arguments: + +- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets. +- `--output_dir`: download to a specified folder. By default, it will download to the current directory. + +Start a server with speculative decoding: + +```bash +vllm serve meta-llama/Llama-3.3-70B-Instruct \ + --speculative-config $'{"method": "eagle3", + "num_speculative_tokens": 3, + "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}' +``` + +Run all categories in the Qualitative split: + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 +``` + +Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`. + +Run only a specific category like "multilingual": + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 + --speed-bench-category "multilingual" +``` + +Run all categories in the Throughput split (2k ISL): + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --speed-bench-dataset-subset throughput_2k + --dataset-path "/data/speed_bench/" \ + --num-prompts -1 +``` + +Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card. + #### Other HuggingFaceDataset Examples ```bash diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index 56329a6edcc5..472e1cf57ff9 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -140,6 +140,80 @@ Data parallelism replicates the entire model across multiple GPU sets and proces Data parallelism can be combined with the other parallelism strategies and is set by `data_parallel_size=N`. Note that MoE layers will be sharded according to the product of the tensor parallel size and data parallel size. +### NUMA Binding for Multi-Socket GPU Nodes + +On multi-socket GPU servers, GPU worker processes can lose performance if their +CPU execution and memory allocation drift away from the NUMA node nearest to the +GPU. vLLM can pin each worker with `numactl` before the Python subprocess starts, +so the interpreter, imports, and early allocator state are created with the +desired NUMA policy from the beginning. + +Use `--numa-bind` to enable the feature. By default, vLLM auto-detects the +GPU-to-NUMA mapping and uses `--cpunodebind= --membind=` for each +worker. When you need a custom CPU policy, add `--numa-bind-cpus` and vLLM will +switch to `--physcpubind= --membind=`. + +These `--numa-bind*` options only apply to GPU execution processes. They do not +configure the CPU backend's separate thread-affinity controls. Automatic +GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as +ROCM-based platforms; other GPU backends must provide explicit binding lists if +they use these options. + +`--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in +the same order as the GPU indices. +`--numa-bind-cpus` takes one `numactl` CPU list per visible GPU, in the same +order as the GPU indices. Each CPU list must use +`numactl --physcpubind` syntax such as `0-3`, `0,2,4-7`, or `16-31,48-63`. + +```bash +# Auto-detect NUMA nodes for visible GPUs +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --tensor-parallel-size 4 \ + --numa-bind + +# Explicit NUMA-node mapping +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --tensor-parallel-size 4 \ + --numa-bind \ + --numa-bind-nodes 0 0 1 1 + +# Explicit CPU pinning, useful for PCT or other high-frequency core layouts +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --tensor-parallel-size 4 \ + --numa-bind \ + --numa-bind-nodes 0 0 1 1 \ + --numa-bind-cpus 0-3 4-7 48-51 52-55 +``` + +Notes: + +- CLI usage forces multiprocessing to use the `spawn` method automatically. If you enable NUMA binding through the Python API, also set `VLLM_WORKER_MULTIPROC_METHOD=spawn`. +- Automatic detection relies on NVML and NUMA support from the host. If it cannot determine the mapping reliably, pass `--numa-bind-nodes` explicitly. +- Explicit `--numa-bind-nodes` and `--numa-bind-cpus` values must be valid `numactl` inputs. vLLM does a small amount of validation, but the effective binding semantics are still determined by `numactl`. +- The current implementation binds GPU execution processes such as `EngineCore` and multiprocessing workers. It does not apply NUMA binding to frontend API server processes or the DP coordinator. +- In containerized environments, NUMA policy syscalls may require extra permissions, such as `--cap-add SYS_NICE` when running via `docker run`. + +### CPU Backend Thread Affinity + +The CPU backend uses a different mechanism from `--numa-bind`. CPU execution is +configured through CPU-specific environment variables such as +`VLLM_CPU_OMP_THREADS_BIND`, `VLLM_CPU_NUM_OF_RESERVED_CPU`, and +`CPU_VISIBLE_MEMORY_NODES`, rather than the GPU-oriented `--numa-bind*` CLI +options. + +By default, `VLLM_CPU_OMP_THREADS_BIND=auto` derives OpenMP placement from the +available CPU and NUMA topology for each CPU worker. To override the automatic +policy, set `VLLM_CPU_OMP_THREADS_BIND` explicitly using the CPU list format +documented for the CPU backend, or use `nobind` to disable this behavior. + +For the current CPU backend setup and tuning guidance, see: + +- [Related runtime environment variables](../getting_started/installation/cpu.md#related-runtime-environment-variables) +- [How to decide `VLLM_CPU_OMP_THREADS_BIND`](../getting_started/installation/cpu.md#how-to-decide-vllm_cpu_omp_threads_bind) + +The GPU-only `--numa-bind`, `--numa-bind-nodes`, and `--numa-bind-cpus` options +do not configure CPU worker affinity. + ### Batch-level DP for Multi-Modal Encoders By default, TP is used to shard the weights of multi-modal encoders just like for language decoders, diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 24e7d1c5be06..a538a408e986 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -49,10 +49,10 @@ If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first: uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 ``` -Then install the necessary build dependencies from `requirements/build.txt`, skipping `torch` as it was installed in the previous step: +Then install the necessary build dependencies from `requirements/build/cuda.txt`, skipping `torch` as it was installed in the previous step: ```bash -grep -v '^torch==' requirements/build.txt | uv pip install -r - +grep -v '^torch==' requirements/build/cuda.txt | uv pip install -r - ``` Finally install vLLM using: diff --git a/docs/contributing/incremental_build.md b/docs/contributing/incremental_build.md index cc01a60ce1e7..6be35af6131b 100644 --- a/docs/contributing/incremental_build.md +++ b/docs/contributing/incremental_build.md @@ -16,10 +16,10 @@ Before setting up the incremental build: 2. **CUDA Toolkit:** Verify that the NVIDIA CUDA Toolkit is correctly installed and `nvcc` is accessible in your `PATH`. CMake relies on `nvcc` to compile CUDA code. You can typically find `nvcc` in `$CUDA_HOME/bin/nvcc` or by running `which nvcc`. If you encounter issues, refer to the [official CUDA Toolkit installation guides](https://developer.nvidia.com/cuda-toolkit-archive) and vLLM's main [GPU installation documentation](../getting_started/installation/gpu.md#troubleshooting) for troubleshooting. The `CMAKE_CUDA_COMPILER` variable in your `CMakeUserPresets.json` should also point to your `nvcc` binary. -3. **Build Tools:** It is highly recommended to install `ccache` for fast rebuilds by caching compilation results (e.g., `sudo apt install ccache` or `conda install ccache`). Also, ensure the core build dependencies like `cmake` and `ninja` are installed. These are installable through `requirements/build.txt` or your system's package manager. +3. **Build Tools:** It is highly recommended to install `ccache` for fast rebuilds by caching compilation results (e.g., `sudo apt install ccache` or `conda install ccache`). Also, ensure the core build dependencies like `cmake` and `ninja` are installed. These are installable through `requirements/build/cuda.txt` or your system's package manager. ```console - uv pip install -r requirements/build.txt --torch-backend=auto + uv pip install -r requirements/build/cuda.txt --torch-backend=auto ``` ## Setting up the CMake Build Environment diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index a23de100da39..3e2ee38d2bdd 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model: See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls. -Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns: +Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns: #### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n) @@ -75,21 +75,20 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt ??? code "get_generation_prompt()" ```python + from vllm.config.speech_to_text import SpeechToTextParams + class YourASRModel(nn.Module, SupportsTranscription): ... @classmethod def get_generation_prompt( cls, - audio: np.ndarray, - stt_config: SpeechToTextConfig, - model_config: ModelConfig, - language: str | None, - task_type: Literal["transcribe", "translate"], - request_prompt: str, - to_language: str | None, + stt_params: SpeechToTextParams, ) -> PromptType: - # Example with a free-form instruction prompt + audio = stt_params.audio + stt_config = stt_params.stt_config + task_type = stt_params.task_type + task_word = "Transcribe" if task_type == "transcribe" else "Translate" prompt = ( "user\n" @@ -112,20 +111,22 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries: ??? code "get_generation_prompt()" ```python + from vllm.config.speech_to_text import SpeechToTextParams + class YourASRModel(nn.Module, SupportsTranscription): ... @classmethod def get_generation_prompt( cls, - audio: np.ndarray, - stt_config: SpeechToTextConfig, - model_config: ModelConfig, - language: str | None, - task_type: Literal["transcribe", "translate"], - request_prompt: str, - to_language: str | None, + stt_params: SpeechToTextParams, ) -> PromptType: + audio = stt_params.audio + stt_config = stt_params.stt_config + language = stt_params.language + task_type = stt_params.task_type + request_prompt = stt_params.request_prompt + if language is None: raise ValueError("Language must be specified") @@ -193,7 +194,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics: The API server takes care of basic audio I/O and optional chunking before building prompts: -- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`. +- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`. - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`. - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words. @@ -206,22 +207,20 @@ Relevant server logic: async def _preprocess_speech_to_text(...): language = self.model_cls.validate_language(request.language) ... - y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate) - duration = librosa.get_duration(y=y, sr=sr) + y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate) + duration = get_audio_duration(y=y, sr=sr) do_split_audio = (self.asr_config.allow_audio_chunking and duration > self.asr_config.max_audio_clip_s) chunks = [y] if not do_split_audio else self._split_audio(y, int(sr)) prompts = [] for chunk in chunks: - prompt = self.model_cls.get_generation_prompt( + stt_params = request.build_stt_params( audio=chunk, stt_config=self.asr_config, model_config=self.model_config, - language=language, task_type=self.task_type, - request_prompt=request.prompt, - to_language=to_language, ) + prompt = self.model_cls.get_generation_prompt(stt_params) prompts.append(prompt) return prompts, duration ``` diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 1d12d63549a0..addda300d020 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -206,8 +206,8 @@ Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_cont used to profile a section of code. !!! note - The legacy import paths `vllm.utils.cprofile` and `vllm.utils.cprofile_context` are deprecated. - Please use `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` instead. + The `vllm.utils.profiling` helpers are deprecated and will be removed in + `v0.21`. Please use Python's `cProfile` module directly instead. ### Example usage - decorator diff --git a/docs/deployment/k8s.md b/docs/deployment/k8s.md index dbcb277278c9..7a92c99b2c4a 100644 --- a/docs/deployment/k8s.md +++ b/docs/deployment/k8s.md @@ -4,6 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le - [Deployment with CPUs](#deployment-with-cpus) - [Deployment with GPUs](#deployment-with-gpus) +- [Serving with gRPC](#serving-with-grpc) - [Troubleshooting](#troubleshooting) - [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated) - [Conclusion](#conclusion) @@ -387,6 +388,49 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) If the service is correctly deployed, you should receive a response from the vLLM model. +## Serving with gRPC + +vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies: + +```bash +pip install vllm[grpc] +``` + +When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24). + +To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes: + +```yaml +containers: +- name: mistral-7b + image: vllm/vllm-openai:latest + command: ["/bin/sh", "-c"] + args: [ + "pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code" + ] + ports: + - containerPort: 50051 + livenessProbe: + grpc: + port: 50051 + initialDelaySeconds: 120 + periodSeconds: 10 + readinessProbe: + grpc: + port: 50051 + initialDelaySeconds: 120 + periodSeconds: 5 +``` + +!!! note + The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`. + +You can also verify the health service manually with `grpcurl`: + +```bash +grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check +``` + ## Troubleshooting ### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated" diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 242cc6b3b1ed..c7d171c165aa 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first). | 2 | `FLASH_ATTN` | | 3 | `TRITON_ATTN` | | 4 | `FLEX_ATTENTION` | +| 5 | `TURBOQUANT` | **Ampere/Hopper (SM 8.x-9.x):** @@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first). | 2 | `FLASHINFER` | | 3 | `TRITON_ATTN` | | 4 | `FLEX_ATTENTION` | +| 5 | `TURBOQUANT` | ### MLA Attention (DeepSeek-style) @@ -167,10 +169,10 @@ Priority is **1 = highest** (tried first). | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ | | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A | | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 | +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 | +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any | | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A | @@ -178,6 +180,7 @@ Priority is **1 = highest** (tried first). | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any | +| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > @@ -210,9 +213,9 @@ configuration. | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | -| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | +| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 5515c91a8b69..e32010232ef0 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g. class BudgetGraphMetadata: token_budget: int max_batch_size: int + max_frames_per_batch: int graph: torch.cuda.CUDAGraph input_buffer: torch.Tensor # e.g. pixel_values metadata_buffers: dict[str, torch.Tensor] # e.g. embeddings, seq metadata @@ -51,6 +52,15 @@ For each graph replay: When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`. +### Video inference support + +Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. + +!!! note + Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture. + + Mixed inputs (image+video) per prompt are also supported now. + ## Model integration via `SupportsEncoderCudaGraph` Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods: @@ -65,12 +75,18 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay. * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay). * `encoder_eager_forward(...)` — fallback eager forward when no graph fits. - -Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`). +* `get_input_modality(...)` - return the modality of the inputs. +* `get_max_frames_per_video()` - return model-specific max frames per video. !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. +**Supported models:** + +| Architecture | Models | CG for Image | CG for Video | +| ------------ | ------ | ------------ | ------------ | +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | + !!! note Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. @@ -80,10 +96,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs: * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level. * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference. -* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. +* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. +* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode). ## Usage guide +### Image inference + Enable encoder CUDA Graphs via `compilation_config`: ```bash @@ -95,7 +114,7 @@ With explicit budgets: ```bash vllm serve Qwen/Qwen3-VL-32B \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` Python example: @@ -107,7 +126,7 @@ compilation_config = { "cudagraph_mm_encoder": True, # Optional: override auto-inferred budgets # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], - # "encoder_cudagraph_max_images_per_batch": 8, + # "encoder_cudagraph_max_vision_items_per_batch": 8, } model = vllm.LLM( @@ -118,6 +137,41 @@ model = vllm.LLM( The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets). +### Video inference + +Enable encoder CUDA Graphs via `compilation_config`: + +```bash +vllm serve Qwen/Qwen3-VL-32B \ + --compilation-config '{"cudagraph_mm_encoder": true}' +``` + +With explicit budgets: + +```bash +vllm serve Qwen/Qwen3-VL-32B \ + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}' +``` + +Python example: + +```python +import vllm + +compilation_config = { + "cudagraph_mm_encoder": True, + # Optional: override auto-inferred budgets + # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], + # "encoder_cudagraph_max_vision_items_per_batch": 8, + # "encoder_cudagraph_max_frames_per_batch": 64, +} + +model = vllm.LLM( + model="Qwen/Qwen3-VL-32B", + compilation_config=compilation_config, +) +``` + ## About the Performance The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details. @@ -140,7 +194,7 @@ vllm bench mm-processor \ --num-prompts 3000 --num-warmups 300 \ --max-model-len 32768 --seed 42 \ --mm-encoder-attn-backend FLASH_ATTN \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` ### Multi-GPU (4x GB200, TP=4, DP=4) @@ -165,5 +219,8 @@ vllm bench mm-processor \ --max-model-len 8192 --seed 42 \ --mm-encoder-attn-backend FLASHINFER \ --tensor-parallel-size 4 --mm-encoder-tp-mode data \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` + +!!! note + Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061). diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 3fe1f769bb79..046c509d4b82 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -21,6 +21,7 @@ or just on the low or high end. | Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` | | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ | | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low | +| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low | | [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always | | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | @@ -30,6 +31,7 @@ or just on the low or high end. | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always | +| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm → single kernel | O1 (ROCm/AITER only) | ~2% | No | Always | ## Support Matrix @@ -40,8 +42,9 @@ The table below lists the quantization schemes supported by each fusion on each | Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm | | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- | | `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — | +| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — | | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | -| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* | +| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | | `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — | | `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | @@ -49,11 +52,15 @@ The table below lists the quantization schemes supported by each fusion on each | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | +| `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 | \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) for per-backend details. +\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires +tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`. + † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today; other architectures support requires setting `PassConfig.sp_min_token_num` explicitly. SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`. @@ -145,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt - `FLASHINFER`: CUDA sm100+ with FlashInfer installed -`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`: +`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant` The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where @@ -184,6 +191,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py) +### MiniMax QK Norm (`fuse_minimax_qk_norm`) + +!!! info + This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold: + the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`), + and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any + optimization level. + +**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the +per-token Q/K variances before applying RMS normalization to Q and K. + +This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion): +`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while +`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models. + +Example: + +```bash +vllm serve MiniMaxAI/MiniMax-M2.5 \ + --tensor-parallel-size 4 \ + --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}' +``` + +**Code locations.** + +- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py) +- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`) +- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py) + ### Sequence Parallelism (`enable_sp`) **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather, @@ -347,6 +383,44 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled. - Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`) +### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`) + +!!! info + ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention. + +!!! note + When the native implementation of `rms_norm` is used (the default on CUDA and + ROCm for now), Inductor's built-in fusion already handles merging these norms + automatically. This explicit pass targets the case where AITER's custom + `rms_norm` op is active, which Inductor cannot fuse on its own. + +**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm +operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER, +reducing kernel launch overhead from 2 launches to 1 per MLA layer. + +```text +# Unfused: +q_c, kv_lora = split(projected, [q_dim, kv_dim]) +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim]) +q_c = rms_norm(q_c, q_weight, eps) +kv_c = rms_norm(kv_c, kv_weight, eps) + +# Fused: +q_c, kv_lora = split(projected, [q_dim, kv_dim]) +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim]) +q_normed, kv_normed = fused_mla_dual_rms_norm( + q_c, q_weight, kv_c, kv_weight, eps1, eps2) +``` + +Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above +when AITER is available. + +**Code locations.** + +- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`) +- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`) +- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442) + ## See Also - [Optimization Levels](optimization_levels.md) — high-level presets that set diff --git a/docs/design/metrics.md b/docs/design/metrics.md index be917c0dc614..0ae420399767 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us ### Grafana Dashboard -vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard. +vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard. The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important: @@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing: - Added by and reinstated by - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces` - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/) -- [User-facing docs](../../examples/online_serving/opentelemetry/README.md) +- [User-facing docs](../../examples/observability/opentelemetry/README.md) - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f) - [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 7cf7b76d6ed4..54b796fde3bf 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -36,7 +36,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | -| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 @@ -59,7 +59,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes. - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod] - [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod] - [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod] -- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod] +- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod] - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod] ## Fused Experts Kernels @@ -82,11 +82,11 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | ------ | ----------------- | ------------ | ------------- | ------------------- | --------------------- | ------- | ------ | | triton | standard | all1 | G,A,T | silu, gelu,
swigluoai,
silu_no_mul,
gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],
[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] | | triton (batched) | batched | all1 | G,A,T | silu, gelu | 6 | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] | -| deep gemm | standard,
batched | fp8 | G(128),A,T | silu, gelu | 6 | Y |
[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],
[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] | -| cutlass_fp4 | standard,
batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] | -| cutlass_fp8 | standard,
batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],
[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] | +| deep gemm | standard,
batched | fp8 | G(128),A,T | silu, gelu | 6 | Y |
[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],
[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] | +| cutlass_fp4 | standard,
batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] | +| cutlass_fp8 | standard,
batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],
[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] | | flashinfer | standard | nvfp4,
fp8 | T | 5 | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] | -| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] | +| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] | | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | diff --git a/docs/design/optimization_levels.md b/docs/design/optimization_levels.md index 591978b542e6..dd0936ca9e5b 100644 --- a/docs/design/optimization_levels.md +++ b/docs/design/optimization_levels.md @@ -56,6 +56,7 @@ Fusions: - `-cc.pass_config.fuse_norm_quant=True`* - `-cc.pass_config.fuse_act_quant=True`* - `-cc.pass_config.fuse_act_padding=True`† +- `-cc.pass_config.fuse_mla_dual_rms_norm=True`† \* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.
† These fusions are ROCm-only and require AITER. diff --git a/docs/design/quota_serve.md b/docs/design/quota_serve.md new file mode 100644 index 000000000000..e0de3cc7e339 --- /dev/null +++ b/docs/design/quota_serve.md @@ -0,0 +1,108 @@ +# QuotaServe: session demand and shared KV cache quotas + +```bash +EVICTION_POLICY=quotaserve QUOTASERVE_APPS='["chat","agent"]' \ + QUOTASERVE_SESSION_IDLE_TTL_S=60 \ + vllm serve --enable-prefix-caching --enable-prompt-tokens-details +``` + +`EVICTION_POLICY` selects `lru` (default), `lfu`, `fifo`, or `quotaserve`. +LRU uses the native free queue. FIFO orders cached blocks by insertion; LFU +counts prefix acquisitions and breaks frequency ties by recency. All policies +reclaim only blocks with no active request references and prefer uncached space. + +QuotaServe requires prefix caching and a JSON list of unique, nonempty app IDs. +Requests supply `application_id` and `session_id` through `vllm_xargs`. +Unregistered apps can borrow capacity but have no protected quota. App cache +salts apply to every policy in the comparison. + +## Session demand + +The scheduler registers sessions when requests arrive, including queued +requests. On normal completion, it counts the request's unique, non-null +physical KV blocks before release. Cached hits and partially filled blocks +count toward this footprint. Aborted requests do not update the footprint mean. +For sliding-window and hybrid models, the observation is the final resident +footprint, not necessarily the peak. + +Each session retains its latest completed footprint while waiting for the next +turn. For full-attention cache groups, the scheduler also estimates blocks from +the arrived prompt, accounting for context-parallel block size. The larger of +that estimate and the latest completed footprint enters the session's demand. +This uses an already received request, without future-turn information. Other +cache types use the completed-footprint estimate. + +Apps maintain a request-footprint EWMA with weight 0.2 for new samples. Sessions +without a footprint or arrived-prompt estimate use this mean. An unobserved app +uses the mean of observed apps; before any observations the prior is usable +capacity divided by registered app count. The app score sums known session +demands and prior estimates for unknown sessions, with at least one prior share +for an app that has no tracked sessions. Scores round up to whole blocks. + +If the total score fits in capacity, quotas equal scores and unused capacity +remains shared. Otherwise, scores scale proportionally and use largest-remainder +rounding. These are soft quotas, not physical reservations or guarantees for +individual session prefixes. + +## Think time and session retirement + +A session becomes idle when its last unfinished request completes. A subsequent +request from a still-tracked idle session updates its app's think-time EWMA and +EWMA variance. With weight `alpha = 0.2`, old mean `G`, variance `V`, and new gap +`g`, the updates are: + +```text +V <- (1 - alpha) * (V + alpha * (g - G)^2) +G <- (1 - alpha) * G + alpha * g +``` + +The first gap initializes the mean and zero variance. Idle timeout is +`max(5 seconds, 3 * G)` once a gap is observed; before that it uses +`QUOTASERVE_SESSION_IDLE_TTL_S`. Queued or running sessions do not expire. +Expiration is checked during policy updates, without a background timer. + +The policy retains block references from each session's latest completed +context. A timeout or a newly completed context can retire older references. +Only valid, unreferenced cached blocks no longer retained by another tracked +session enter the retired-block LRU queue. A new acquisition removes the block +from eviction candidates. + +## Block allocation and eviction + +Resident occupancy includes both pinned blocks and unreferenced cached blocks. +A cache hit does not reduce occupancy. Uncached blocks stop counting when their +last request reference is released. Allocation proceeds one block at a time: + +1. Use uncached free space. +2. Reclaim the oldest retired cached block, if available + (`selection="quotaserve_retired"`). +3. Find owners with evictable cache and + `resident[app] + (app == requester) > quota[app]`. Choose the oldest LRU head + across these owners, independently of the amount of quota excess. +4. Before reclaiming that candidate, compare its time since release with + `min(idle_timeout, G + 2 * sqrt(V))` for its owner. Within this reuse interval, + use global LRU instead (`quotaserve_reuse`). Otherwise reclaim the quota + candidate (`quotaserve`). Without a gap observation the interval is zero. +5. If no owner has an eligible quota-excess block, use global LRU + (`quotaserve_pressure`) so requests can continue. + +The reuse interval is a heuristic based on observed gaps, not a statistical +coverage guarantee. Both global-LRU paths can reclaim cache within another +app's quota. Neither path changes request references or permits reclaiming +pinned blocks. No app-specific think time is hardcoded. + +## Validation and GPU replay + +```bash +.venv/bin/python -m pytest tests/v1/core/test_quota_serve.py \ + tests/v1/core/test_ranked_eviction.py -q +``` + +The sibling research repository contains +`004_serving_experiments/runner.py`. Its GPU configuration uses +Llama-3.2-3B-Instruct on an A10G, fixed KV memory, and a 64-request execution +limit. The saved input has 60 Chat sessions and 120 Agent sessions, each with +10 original turns. Chat waits 20 seconds and Agent waits 1 second after each +response. Both reuse original contexts and generate the original response's +token length. See that directory's README and `results/gpu_comparison.md` for +conditions, raw measurements, repeat counts, and limitations. diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 85487697fd37..804cd905e3b1 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -104,7 +104,7 @@ for output in outputs: Batch invariance has been tested and verified on the following models: - **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1` -- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B` +- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ` - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct` - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` - **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct` diff --git a/docs/features/context_extension.md b/docs/features/context_extension.md new file mode 100644 index 000000000000..f622191aebc6 --- /dev/null +++ b/docs/features/context_extension.md @@ -0,0 +1,70 @@ +# Context Extension + +!!! note + The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead. +This directory contains examples for extending the context length of models using vLLM. + +## Offline Inference Example + +The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example. + +### Usage + +```bash +python examples/offline_inference/context_extension.py +``` + +## OpenAI Online Method + +You can also use vLLM's OpenAI-compatible API to serve models with extended context length. + +### Usage + +Run the vLLM server with the following command to extend the context length using YARN: + +```bash +vllm serve Qwen/Qwen3-0.6B \ + --hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \ + --max-model-len 131072 +``` + +### Client Example + +After starting the server, you can use the OpenAI Python client to interact with it: + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8000/v1", + api_key="token-abc123" # Dummy API key, required by the client +) + +response = client.chat.completions.create( + model="Qwen/Qwen3-0.6B", + messages=[ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "Hello"} + ], + max_tokens=128, + temperature=0.8, + top_p=0.95 +) + +print(response.choices[0].message.content) +``` + +### Key Parameters + +The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters). + +Common parameters include: + +- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic") +- `factor`: The factor by which to extend the context length +- `original_max_position_embeddings`: The original maximum position embeddings of the model + +The following parameters are specific to vLLM: + +- `max_model_len`: The new maximum sequence length after extension (original * factor). + Used for KV cache pre‑allocation and request limit at serving time. diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index d95427464196..af6da94aa622 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -72,4 +72,4 @@ For the PD disaggregation part, the Prefill instance receives cache exactly the `docs/features/disagg_prefill.md` shows the brief idea about the disaggregated prefill (v0) -We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D; +We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl/` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D; diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index ee82c34fa0eb..df7aef3f14ff 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech. ```python -import librosa from vllm import LLM, SamplingParams from vllm.multimodal.audio import split_audio +from vllm.multimodal.media.audio import load_audio # Load long audio file -audio, sr = librosa.load("long_audio.wav", sr=16000) +audio, sr = load_audio("long_audio.wav", sr=16000) # Split into chunks at low-energy (quiet) regions chunks = split_audio( @@ -780,6 +780,70 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ Works with common video formats like MP4 when using OpenCV backends. +#### Pre-extracted Frame Sequences with `media_io_kwargs` + +When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction. + +**Supported Parameters:** + +| Parameter | Type | Description | +| --------- | ---- | ----------- | +| `fps` | float | Frame rate of the original video | +| `frames_indices` | list[int] | Indices of the actually sampled frames | +| `total_num_frames` | int | Total frame count of the original video | +| `duration` | float | Duration of the original video in seconds | +| `do_sample_frames` | bool | Whether to perform frame sampling | + +??? code + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") + + # Client-side frame extraction + frames = extract_frames(video_path, num_frames=32) + frames_b64 = ",".join([encode_image(f) for f in frames]) + video_url = f"data:video/jpeg;base64,{frames_b64}" + + # Pass video metadata via media_io_kwargs + response = client.chat.completions.create( + model="your-multimodal-model", + messages=[{ + "role": "user", + "content": [ + {"type": "video_url", "video_url": {"url": video_url}}, + {"type": "text", "text": "Describe what happens in this video."} + ] + }], + extra_body={ + "media_io_kwargs": { + "video": { + "fps": 30.0, + "frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, + 100, 110, 120, 130, 140, 150, 160, 170, + 180, 190, 200, 210, 220, 230, 240, 250, + 260, 270, 280, 290, 300, 310], + "total_num_frames": 900, + "duration": 30.0, + } + } + }, + ) + + print(response.choices[0].message.content) + ``` + +**Why use `media_io_kwargs`?** + +When extracting frames client-side, the server loses important context about the original video: + +- **Temporal information**: Which frames were sampled and their positions in the original timeline +- **Video duration**: How long the original video was +- **Frame rate**: The original playback speed + +By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped. + #### Custom RGBA Background Color To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`: @@ -832,7 +896,7 @@ Then, you can use the OpenAI client as follows: base_url=openai_api_base, ) - # Any format supported by librosa is supported + # Any format supported by soundfile/PyAV is supported audio_url = AudioAsset("winning_call").url audio_base64 = encode_base64_content_from_url(audio_url) diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 0b8fc71d3f30..4be088c56a4c 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -16,6 +16,7 @@ The following are the supported quantization formats for vLLM: - [INT8 W8A8](int8.md) - [FP8 W8A8](fp8.md) - [NVIDIA Model Optimizer](modelopt.md) +- [Online Quantization](online.md) - [AMD Quark](quark.md) - [Quantized KV Cache](quantized_kvcache.md) - [TorchAO](torchao.md) diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md new file mode 100644 index 000000000000..913da250ee67 --- /dev/null +++ b/docs/features/quantization/online.md @@ -0,0 +1,97 @@ +# Online Quantization + +Online quantization lets you take a BF16/FP16 model and quantize its Linear +and MoE weights to lower precision (such as FP8) at load time, without needing +a pre-quantized checkpoint or calibration data. Weights are converted during +model loading and activations are dynamically scaled during each forward pass. + +## Quick Start + +Pass a scheme name to the `quantization` parameter: + +```python +from vllm import LLM + +# Per-tensor FP8 quantization (one scale per weight tensor) +llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor") + +# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations) +llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block") + +# MXFP8 quantization for weights and activations +llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8") +``` + +Or with the CLI: + +```bash +vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor +vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block +vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8 +``` + +## Supported Schemes + +| Scheme | Weight recipe | Activation recipe | Notes | +| ------ | ------------- | ------------------ | ----- | +| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | +| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | +| `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback | + +## Advanced Configuration + +For fine-grained control, use a `quantization_config` dictionary. + +### Separate Schemes for Dense and MoE Layers + +You can apply different quantization schemes to dense linear layers and MoE expert layers: + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "linear_scheme_override": "fp8_per_block", + }, +) +``` + +Or, + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "moe_scheme_override": "fp8_per_block", + }, +) +``` + +### Excluding Layers from Quantization + +Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`): + +```python +from vllm import LLM + +llm = LLM( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="fp8_per_tensor", + quantization_config={ + "ignore": [ + # exact layer name + "model.layers.1.self_attn.o_proj", + # regex: skip all QKV projections + "re:.*[qkv]_proj", + ], + }, +) +``` + +!!! note + For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name. diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 9d41185f1ae1..c7b2d688f228 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -249,7 +249,7 @@ Token counting starts from `reasoning_start_str`. Once the reasoning token count To use this feature: - `--reasoning-parser` enables reasoning extraction. -- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`). +- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`). If not set, vLLM will attempt to automatically initialize these tokens from the reasoning parser. - `thinking_token_budget` (a sampling parameter) sets the per-request reasoning token limit. If `thinking_token_budget` is not specified, no explicit reasoning limit is applied beyond normal generation constraints such as `max_tokens`. @@ -283,9 +283,7 @@ curl http://localhost:8000/v1/chat/completions \ "messages": [ { "role": "user", "content": "9.11 and 9.8, which is greater?" } ], - "extra_body": { - "thinking_token_budget": 10 - } + "thinking_token_budget": 10 }' ``` diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 9793de3f4c35..25cda8059b24 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -35,6 +35,99 @@ For reproducible measurements in your environment, use [`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py) or the [benchmark CLI guide](../../benchmarking/cli.md). +## `--speculative-config` schema + +Use `--speculative-config` to pass speculative decoding settings as a JSON +object on the CLI: + +```bash +vllm serve \ + --speculative-config '{ + "method": "draft_model", + "model": "", + "num_speculative_tokens": 5 + }' +``` + +The same keys are accepted from Python via `LLM(..., speculative_config={...})`. +The tables below highlight common user-facing keys accepted in this JSON +object; they are not an exhaustive schema reference. +For more details, see the generated [engine arguments reference](../../configuration/engine_args.md) +and the API docs for [vllm.config.SpeculativeConfig][]. + +### Common keys + +These keys are commonly used across speculative decoding setups, though some +only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and +`dflash`. + +| Key | Type | Default | Allowed values / meaning | +| --- | --- | --- | --- | +| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. | +| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. | +| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. | +| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. | +| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. | +| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. | +| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. | +| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. | + +### Method-specific keys + +#### N-gram + +| Key | Type | Default | Meaning | +| --- | --- | --- | --- | +| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. | +| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. | + +Example: + +```bash +vllm serve \ + --speculative-config '{ + "method": "ngram", + "num_speculative_tokens": 4, + "prompt_lookup_min": 2, + "prompt_lookup_max": 5 + }' +``` + +#### Suffix decoding + +| Key | Type | Default | Meaning | +| --- | --- | --- | --- | +| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. | +| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. | +| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. | +| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. | + +Example: + +```bash +vllm serve \ + --speculative-config '{ + "method": "suffix", + "num_speculative_tokens": 8, + "suffix_decoding_max_tree_depth": 24, + "suffix_decoding_max_cached_requests": 10000, + "suffix_decoding_max_spec_factor": 1.0, + "suffix_decoding_min_token_prob": 0.1 + }' +``` + +### Notes + +- `--speculative-config` expects a JSON object on the CLI. In YAML config + files, use a nested mapping instead of an escaped JSON string. +- `tensor_parallel_size` is not a valid key in `speculative_config`. Use + `draft_tensor_parallel_size` instead. +- Keys such as `temperature` and `top_p` are sampling parameters, not + `--speculative-config` fields. +- Internal fields such as `target_model_config`, `draft_model_config`, + `target_parallel_config`, `draft_parallel_config`, and `draft_load_config` + are populated by vLLM and are not intended to be set by users. + ## Lossless guarantees of Speculative Decoding In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of diff --git a/docs/features/speculative_decoding/draft_model.md b/docs/features/speculative_decoding/draft_model.md index ee0eaf176e76..b4662e6438f9 100644 --- a/docs/features/speculative_decoding/draft_model.md +++ b/docs/features/speculative_decoding/draft_model.md @@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \ --port 8000 \ --seed 42 \ -tp 1 \ - --max_model_len 2048 \ - --gpu_memory_utilization 0.8 \ - --speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}' + --max-model-len 2048 \ + --gpu-memory-utilization 0.8 \ + --speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}' ``` The code used to request as completions as a client remains unchanged: @@ -77,4 +77,8 @@ The code used to request as completions as a client remains unchanged: ``` !!! warning - Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated. + Note: Please use `--speculative-config` to set all configurations related + to speculative decoding. The previous method of specifying the model + through `--speculative-model` and adding related parameters such as + `--num-speculative-tokens` separately has been deprecated. For supported + keys and examples, see the [`--speculative-config` schema](README.md#--speculative-config-schema). diff --git a/docs/features/speculative_decoding/mtp.md b/docs/features/speculative_decoding/mtp.md index bcd7153deb51..7e1d1ec7038b 100644 --- a/docs/features/speculative_decoding/mtp.md +++ b/docs/features/speculative_decoding/mtp.md @@ -38,7 +38,7 @@ for output in outputs: ```bash vllm serve XiaomiMiMo/MiMo-7B-Base \ --tensor-parallel-size 1 \ - --speculative_config '{"method":"mtp","num_speculative_tokens":1}' + --speculative-config '{"method":"mtp","num_speculative_tokens":1}' ``` ## Notes diff --git a/docs/features/speculative_decoding/parallel_draft_model.md b/docs/features/speculative_decoding/parallel_draft_model.md index 2a3f11a302d3..c31b8e2d2f40 100644 --- a/docs/features/speculative_decoding/parallel_draft_model.md +++ b/docs/features/speculative_decoding/parallel_draft_model.md @@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \ --port 8000 \ --seed 42 \ -tp 1 \ - --max_model_len 2048 \ - --gpu_memory_utilization 0.8 \ - --speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}' + --max-model-len 2048 \ + --gpu-memory-utilization 0.8 \ + --speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}' ``` ## Pre-trained PARD weights diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index b266e96db559..8d9cc65f2be2 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -96,14 +96,14 @@ cd vllm_source Third, install required dependencies: ```bash -uv pip install -r requirements/cpu-build.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu uv pip install -r requirements/cpu.txt --torch-backend cpu ``` ??? console "pip" ```bash pip install --upgrade pip - pip install -v -r requirements/cpu-build.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install -v -r requirements/build/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu pip install -v -r requirements/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu ``` diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index eeb20b8bf063..1e36b4317647 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -3,15 +3,15 @@ vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform. -Currently, the CPU implementation for s390x architecture supports FP32 datatype only. +Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16. --8<-- [end:installation] --8<-- [start:requirements] - OS: `Linux` -- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools +- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools - Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above. -- Build install python packages: `pyarrow`, `torch` and `torchvision` +- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless` --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -24,13 +24,14 @@ Currently, there are no pre-built IBM Z CPU wheels. --8<-- [end:pre-built-wheels] --8<-- [start:build-wheel-from-source] -Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4: +Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6: ```bash dnf install -y \ - which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \ + which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \ libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \ - openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel + openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \ + clang llvm-devel llvm-static clang-devel ``` Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation. @@ -43,13 +44,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \ Execute the following commands to build and install vLLM from source. !!! tip - Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM. + Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM. ```bash - sed -i '/^torch/d' requirements/build.txt # remove torch from requirements/build.txt since we use nightly builds uv pip install -v \ + --extra-index-url https://download.pytorch.org/whl/cpu \ --torch-backend auto \ - -r requirements/build.txt \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv pip install dist/*.whl @@ -57,10 +58,9 @@ Execute the following commands to build and install vLLM from source. ??? console "pip" ```bash - sed -i '/^torch/d' requirements/build.txt # remove torch from requirements/build.txt since we use nightly builds pip install -v \ - --extra-index-url https://download.pytorch.org/whl/nightly/cpu \ - -r requirements/build.txt \ + --extra-index-url https://download.pytorch.org/whl/cpu \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ pip install dist/*.whl diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index 8b855e919f44..ad051d22dc8c 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -88,14 +88,14 @@ cd vllm_source Install the required dependencies: ```bash -uv pip install -r requirements/cpu-build.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu uv pip install -r requirements/cpu.txt --torch-backend cpu ``` ??? console "pip" ```bash pip install --upgrade pip - pip install -v -r requirements/cpu-build.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install -v -r requirements/build/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu pip install -v -r requirements/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu ``` diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index e46fecc45cd5..309dd671251e 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -1,12 +1,12 @@ --8<-- [start:installation] -vLLM contains pre-compiled C++ and CUDA (12.8) binaries. +vLLM contains pre-compiled C++ and CUDA (12.9) binaries. --8<-- [end:installation] --8<-- [start:requirements] -- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.) +- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, B200, etc.) --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -31,7 +31,7 @@ uv pip install vllm --torch-backend=auto pip install vllm --extra-index-url https://download.pytorch.org/whl/cu129 ``` -We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu128`), set `--torch-backend=cu128` (or `UV_TORCH_BACKEND=cu128`). If this doesn't work, try running `uv self update` to update `uv` first. +We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu130`), set `--torch-backend=cu130` (or `UV_TORCH_BACKEND=cu130`). If this doesn't work, try running `uv self update` to update `uv` first. !!! note NVIDIA Blackwell GPUs (B200, GB200) require a minimum of CUDA 12.8, so make sure you are installing PyTorch wheels with at least that version. PyTorch itself offers a [dedicated interface](https://pytorch.org/get-started/locally/) to determine the appropriate pip command to run for a given target configuration. @@ -93,7 +93,7 @@ If you only need to change Python code, you can build and install vLLM without c ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -VLLM_USE_PRECOMPILED=1 uv pip install --editable . +VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto ``` This command will do the following: @@ -107,10 +107,10 @@ This command will do the following: 1. If you change C++ or kernel code, you cannot use Python-only build; otherwise you will see an import error about library not found or undefined symbol. 2. If you rebase your dev branch, it is recommended to uninstall vllm and re-run the above command to make sure your libraries are up to date. -In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the main branch was just merged and the wheel is being built. In this case, you can wait for around an hour to try again, or manually assign the previous commit in the installation using the `VLLM_PRECOMPILED_WHEEL_LOCATION` environment variable. +In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`. ```bash -export VLLM_PRECOMPILED_WHEEL_COMMIT=$(git rev-parse HEAD~1) # or earlier commit on main +export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly export VLLM_USE_PRECOMPILED=1 uv pip install --editable . ``` @@ -134,7 +134,7 @@ If you want to modify C++ or CUDA code, you'll need to build vLLM from source. T ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -uv pip install -e . +uv pip install -e . --torch-backend=auto ``` !!! tip @@ -162,7 +162,7 @@ To build vLLM using an existing PyTorch installation: git clone https://github.com/vllm-project/vllm.git cd vllm python use_existing_torch.py -uv pip install -r requirements/build.txt +uv pip install -r requirements/build/cuda.txt uv pip install --no-build-isolation -e . ``` @@ -185,7 +185,7 @@ To achieve this, you can set the environment variable VLLM_CUTLASS_SRC_DIR to po ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -VLLM_CUTLASS_SRC_DIR=/path/to/cutlass uv pip install -e . +VLLM_CUTLASS_SRC_DIR=/path/to/cutlass uv pip install -e . --torch-backend=auto ``` ##### Troubleshooting @@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command. ```bash DOCKER_BUILDKIT=1 docker build \ - --build-arg CUDA_VERSION=13.0.1 \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \ + --build-arg CUDA_VERSION=13.0.2 \ + --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ --build-arg max_jobs=256 \ --build-arg nvcc_threads=2 \ --build-arg RUN_WHEEL_CHECK=false \ diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index 4ab01ee8c687..f8385997eea3 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \ # Install dependencies pip install --upgrade numba \ scipy \ - huggingface-hub[cli,hf_transfer] \ + huggingface-hub[cli] \ setuptools_scm pip install -r requirements/rocm.txt diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index 9e71860d62fd..8c282582281e 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt !!! note - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues. - - For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu). + - For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu). - Finally, build and install vLLM XPU backend: diff --git a/docs/governance/committers.md b/docs/governance/committers.md index df874418f1c4..cff041866b29 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle: - [@aarnphm](https://github.com/aarnphm): Structured output - [@alexm-redhat](https://github.com/alexm-redhat): Performance - [@ApostaC](https://github.com/ApostaC): Connectors, offloading +- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser - [@benchislett](https://github.com/benchislett): Engine core and spec decode - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser @@ -31,6 +32,7 @@ Sorted alphabetically by GitHub handle: - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance - [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed - [@markmc](https://github.com/markmc): Observability +- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance - [@mgoin](https://github.com/mgoin): Quantization and performance - [@NickLucche](https://github.com/NickLucche): KV connector - [@njhill](https://github.com/njhill): Distributed, API server, engine core @@ -41,6 +43,7 @@ Sorted alphabetically by GitHub handle: - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support - [@russellb](https://github.com/russellb): Structured output, engine core, security +- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community - [@tdoublep](https://github.com/tdoublep): State space models @@ -56,6 +59,7 @@ Sorted alphabetically by GitHub handle: - [@zhuohan123](https://github.com/zhuohan123): Project lead, RL integration, numerics - [@zou3519](https://github.com/zou3519): Compilation - [@BoyuanFeng](https://github.com/BoyuanFeng): Compilation, CUDAGraph +- [@xuechendi](https://github.com/xuechendi): Intel CPU/XPU integration, KV connector ### Emeritus Committers @@ -85,7 +89,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints - @robertgshaw2-redhat, @njhill, @russellb - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation - - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg + - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni - KV Connector: Connector interface and implementation for KV cache offload and transfer - @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC - Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks @@ -104,7 +108,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms - @ProExpertProg - Attention: Attention interface for paged attention - - @WoosukKwon, @LucasWilkinson, @heheda12345 + - @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni - FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB - @tlrmchlsmth - Quantization: Various quantization config, weight loading, and kernel. @@ -118,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - State space models: The state space models implementation in vLLM - @tdoublep, @tlrmchlsmth - Reasoning and tool calling parsers - - @chaunceyjiang, @aarnphm + - @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning ### Entrypoints @@ -132,7 +136,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r ### Features - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP. - - @WoosukKwon, @benchislett, @luccafong + - @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni - Structured Output: The structured output implementation - @russellb, @aarnphm - RL: The RL related features such as collective rpc, sleep mode, etc. @@ -152,8 +156,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r ### External Kernels Integration -- FlashAttention: @LucasWilkinson -- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon +- FlashAttention: @LucasWilkinson, @MatthewBonanni +- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni - Blackwell Kernels: @mgoin, @yewentao256 - DeepEP/DeepGEMM: @mgoin, @yewentao256 @@ -175,7 +179,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Plugin Interface: @youkaichao, @Yikun - NVIDIA GPU: @pavanimajety - AMD GPU: @gshtras, @tjtanaa -- Intel CPU/GPU: @jikunshang, @bigPYJ1151 +- Intel CPU/GPU: @jikunshang, @bigPYJ1151, @xuechendi - Google TPU: @yaochengji ### Ecosystem Projects diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 298eea96de3d..a798b8e34530 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -46,7 +46,7 @@ def decorator(cls): # Mock any version checks by reading from compiled CI requirements -with open(ROOT_DIR / "requirements/test.txt") as f: +with open(ROOT_DIR / "requirements/test/cuda.txt") as f: VERSIONS = dict(line.strip().split("==") for line in f if "==" in line) importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0" @@ -59,7 +59,7 @@ class PydanticMagicMock(MagicMock): """`MagicMock` that's able to generate pydantic-core schemas.""" def __init__(self, *args, **kwargs): - name = kwargs.pop("name", None) + name = kwargs.get("name") super().__init__(*args, **kwargs) self.__spec__ = ModuleSpec(name, None) @@ -85,7 +85,8 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): logger.info("Mocking %s for argparse doc generation", e.name) sys.modules[e.name] = PydanticMagicMock(name=e.name) except Exception: - logger.exception("Failed to import %s.%s: %s", module_name, attr) + logger.exception("Failed to import %s.%s", module_name, attr) + raise raise ImportError( f"Failed to import {module_name}.{attr} after mocking {max_mocks} imports" diff --git a/docs/mkdocs/hooks/generate_metrics.py b/docs/mkdocs/hooks/generate_metrics.py index 4565861c4f7f..97282aaee7d0 100644 --- a/docs/mkdocs/hooks/generate_metrics.py +++ b/docs/mkdocs/hooks/generate_metrics.py @@ -19,7 +19,7 @@ "output": "spec_decode.inc.md", }, { - "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py", + "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py", "output": "nixl_connector.inc.md", }, {"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"}, diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index 2cf721f5eefe..c3309c53ca20 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models are a subset of classification models that accept two prompts as input and output num_labels equal to 1. +### Pooling Types + +| Pooling Tasks | Granularity | Description | +|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `CLS` pooling | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. | +| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. | +| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. | +| `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. | +| `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. | + ### Score Types The scoring models is designed to compute similarity scores between two input prompts. It supports three model types @@ -68,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr |-----------------------|---------------|----------------------------------------------|--------------------|--------------------------| | `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier | | `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity | -| `token_classify` | Token-wise | probability vector of classes for each token | nan | nan | +| `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A | | `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) | !!! note @@ -76,14 +86,15 @@ The scoring models is designed to compute similarity scores between two input pr ### Pooling Usages -| Pooling Usages | Description | -|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------| -| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. | -| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). | -| Token Classification Usages | Token-wise classification | -| Token Embedding Usages | Token-wise embedding | -| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. | -| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. | +| Pooling Usages | Description | +|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------| +| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. | +| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). | +| Token Classification Usages | Token-wise classification | +| Token Embedding Usages | Token-wise embedding | +| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. | +| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. | +| Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). | We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs. @@ -91,9 +102,9 @@ For more detailed information, please refer to the link below. - [Classification Usages](classify.md) - [Embedding Usages](embed.md) -- [Reward Usages](reward.md) - [Token Classification Usages](token_classify.md) - [Token Embedding Usages](token_embed.md) +- [Reward Usages](reward.md) - [Scoring Usages](scoring.md) - [Specific Model Examples](specific_models.md) @@ -103,15 +114,17 @@ Each pooling model in vLLM supports one or more of these tasks according to [Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks], enabling the corresponding APIs. -### Offline APIs corresponding to pooling tasks +### Offline APIs corresponding to pooling usages -| Task | APIs | -|------------------|---------------------------------------------------------------------------------------| -| `embed` | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) | -| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)` | -| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` | -| `token_embed` | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)` | -| `plugin` | `LLM.encode(..., pooling_task="plugin")` | +| Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function | +|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------| +| Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier | +| Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity | +| Token Classification Usages | N/A | `token_classify` | N/A | N/A | +| Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) | +| Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A | +| Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A | +| Plugins Usages | N/A | `plugin` | N/A | N/A | !!! note Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled. @@ -137,7 +150,7 @@ It is primarily designed for [score models](scoring.md). The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM. -Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks). +Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages). ### Examples @@ -173,9 +186,12 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats. -Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks). +Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages). + +Code examples: -Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py) +- [Online example](../../../examples/pooling/reward/token_reward_online.py) +- [Offline example](../../../examples/pooling/reward/token_reward_offline.py) ### Examples diff --git a/docs/models/pooling_models/classify.md b/docs/models/pooling_models/classify.md index 69a6fe75d374..c857e4e061ee 100644 --- a/docs/models/pooling_models/classify.md +++ b/docs/models/pooling_models/classify.md @@ -267,12 +267,39 @@ You can modify the `problem_type` via problem_type in the Hugging Face config. T Implement alignment with transformers [ForSequenceClassificationLoss](https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92). -### Logit bias +### Affine Score Calibration -You can modify the `logit_bias` (aka `sigmoid_normalize`) through the logit_bias parameter in `vllm.config.PoolerConfig`. +Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org/wiki/Platt_scaling) (Platt, 1999), is the most widely used method for calibrating classifier outputs into well-calibrated probabilities. + +The calibration follows the transformation: + +`activation((logit - logit_mean) / logit_sigma)` + +| Parameter | Default | Description | +| --------- | ------- | ----------- | +| `logit_mean` | `None` | Mean subtracted from logits (centers scores) | +| `logit_sigma` | `None` | Standard deviation used to scale logits after mean subtraction | + +The computation order is as follows: + +```python +logits -= logit_mean # subtract mean (center scores) +logits /= logit_sigma # divide by sigma (scale) +logits = activation(logits) # e.g. sigmoid +``` + +Example configuration: + +```bash +--pooler-config '{"use_activation": true, "logit_mean": 4.5, "logit_sigma": 1.0}' +``` ## Removed Features ### Remove softmax from PoolingParams We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function. + +### Remove `logit_bias` and `logit_scale` + +`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21. diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 8b3632a9f33c..2f5d1a3fbe09 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | | `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | | `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | +| `JinaEmbeddingsV5Model`C | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ | | `LlamaBidirectionalModel`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ | | `LlamaModel`C, `LlamaForCausalLM`C, `MistralModel`C, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | @@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th !!! note `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights. +!!! note + `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters + (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the + selected adapter into the base weights at load time. Choose the task with + `--hf-overrides '{"jina_task": ""}'`; the default is `retrieval`. + ### Multimodal Models !!! note diff --git a/docs/models/pooling_models/reward.md b/docs/models/pooling_models/reward.md index 8555060e66be..7cb6e8b5bb65 100644 --- a/docs/models/pooling_models/reward.md +++ b/docs/models/pooling_models/reward.md @@ -134,3 +134,13 @@ print(f"Data: {data!r}") ## Online Serving Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary). + +## More examples + +More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward) + +## Deprecated Features + +### `LLM.reward` + +`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead. diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index ac94a0cd76bc..a4424642cd2a 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -160,6 +160,8 @@ The following Score API parameters are supported: --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params" ``` #### Examples @@ -370,6 +372,8 @@ The following rerank api parameters are supported: --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params" --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params" +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params" ``` #### Examples diff --git a/docs/models/pooling_models/token_classify.md b/docs/models/pooling_models/token_classify.md index 201ce4ea6dcb..945a710df96b 100644 --- a/docs/models/pooling_models/token_classify.md +++ b/docs/models/pooling_models/token_classify.md @@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`. Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py). -### As Reward Models +### Reward Models Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md). diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 3396f4eac2d2..c20ae7cf1646 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -9,14 +9,14 @@ - Online APIs: - Pooling API (`/pooling`) -The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token. +The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token. Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md). !!! note Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not - what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or + what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or `--pooler-config.task token_embed` online. ## Typical Use Cases @@ -71,6 +71,14 @@ Models of any architecture can be converted into embedding models using `--conve If your model is not in the above list, we will try to automatically convert the model using [as_embedding_model][vllm.model_executor.models.adapters.as_embedding_model]. +### Special models + +| Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | +| ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `JinaForRanking` | Qwen3-based | `jinaai/jina-reranker-v3` | | | + +jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py) + --8<-- [end:supported-token-embed-models] ## Offline Inference diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b0b6c080ac49..7c87359afb27 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -384,6 +384,7 @@ th { | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ | | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ | | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ | +| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | | | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ | | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ | | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ | @@ -400,6 +401,7 @@ th { | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | +| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ | | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ | | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ | @@ -418,6 +420,7 @@ th { | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | +| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ | | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ | | `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ | | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ | @@ -471,6 +474,7 @@ th { | `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ | | `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ | | `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ | +| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | | | `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ | | `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ | | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ | @@ -550,13 +554,16 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I+ | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ | | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + IE+ | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ | +| `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ | | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | | `Gemma3ForConditionalGeneration` | Gemma 3 | T + IE+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | +| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I+ + V + A* | `google/gemma-4-E2B-it`, etc. | | ✅︎ | | `GLM4VForCausalLM`^ | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | | `GlmOcrForConditionalGeneration` | GLM-OCR | T + IE+ | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ | +| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + IE+ | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ | | `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I+ + V+ | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | | | `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I+ + V+ | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | | @@ -632,16 +639,22 @@ Some models are supported only via the [Transformers modeling backend](#transfor ^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
E Pre-computed embeddings can be inputted for this modality.
+ Multiple items can be inputted per text prompt for this modality. +* Only specific variants of the model support this modality (see notes below).
!!! note `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its MobileNet-v5 vision backbone. - + Performance is not yet fully optimized mainly due to: - - - Both audio and vision MM encoders use `transformers.AutoModel` implementation. + + - Both audio and vision MM encoders use `transformers.AutoModel` implementation. - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups. +!!! note + For `Gemma4ForConditionalGeneration`: + - audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants. + - The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model. + !!! note For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently. @@ -660,11 +673,12 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `CohereAsrForConditionalGeneration` | Cohere-Transcribe | `CohereLabs/cohere-transcribe-03-2026` | | | | `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | | +| `FireRedLIDForConditionalGeneration` | FireRedLID | `PatchyTisa/FireRedLID-vllm`, etc. | | | | `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | | | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | -| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | | ✅︎ | +| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | | @@ -672,6 +686,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition. !!! note `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed. +#### Realtime Transcription + +Speech models that support streaming transcription via the +[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api) +WebSocket endpoint. + +| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | +| ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | | +| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | | + +!!! note + `VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`. + + `Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`. + You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'` + when serving. + ## Pooling Models See [this page](pooling_models/README.md) for more information on how to use pooling models. diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index d75ae7feb49e..fef4df770fa3 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | | `policy` | The policy type for expert parallel load balancing | `"default"` | +| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` | For example: diff --git a/docs/serving/openai_compatible_server.md b/docs/serving/openai_compatible_server.md index 28e643ea5f07..a2c90e3abd49 100644 --- a/docs/serving/openai_compatible_server.md +++ b/docs/serving/openai_compatible_server.md @@ -60,7 +60,7 @@ We currently support the following OpenAI APIs: - [Translation API](#translations-api) (`/v1/audio/translations`) - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription). - [Realtime API](#realtime-api) (`/v1/realtime`) - - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription). + - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription). In addition, we have the following custom APIs: @@ -467,28 +467,11 @@ It consists of two endpoints: - `/tokenize` corresponds to calling `tokenizer.encode()`. - `/detokenize` corresponds to calling `tokenizer.decode()`. -### Score API - -#### Score Template - -Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)). - -Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template. - -Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter: - -- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}` -- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}` - -This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future. - -Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja) - ### Generative Scoring API The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes". -This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models. +This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models. **Requirements:** diff --git a/examples/online_serving/dashboards/README.md b/examples/observability/dashboards/README.md similarity index 92% rename from examples/online_serving/dashboards/README.md rename to examples/observability/dashboards/README.md index 10b9a864f572..e5f5010a42a4 100644 --- a/examples/online_serving/dashboards/README.md +++ b/examples/observability/dashboards/README.md @@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml For detailed deployment instructions and platform-specific options, see: -- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import -- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping +- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import +- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping ## Contributing diff --git a/examples/online_serving/dashboards/grafana/README.md b/examples/observability/dashboards/grafana/README.md similarity index 100% rename from examples/online_serving/dashboards/grafana/README.md rename to examples/observability/dashboards/grafana/README.md diff --git a/examples/online_serving/dashboards/grafana/performance_statistics.json b/examples/observability/dashboards/grafana/performance_statistics.json similarity index 100% rename from examples/online_serving/dashboards/grafana/performance_statistics.json rename to examples/observability/dashboards/grafana/performance_statistics.json diff --git a/examples/online_serving/dashboards/grafana/query_statistics.json b/examples/observability/dashboards/grafana/query_statistics.json similarity index 100% rename from examples/online_serving/dashboards/grafana/query_statistics.json rename to examples/observability/dashboards/grafana/query_statistics.json diff --git a/examples/online_serving/dashboards/perses/README.md b/examples/observability/dashboards/perses/README.md similarity index 100% rename from examples/online_serving/dashboards/perses/README.md rename to examples/observability/dashboards/perses/README.md diff --git a/examples/online_serving/dashboards/perses/performance_statistics.yaml b/examples/observability/dashboards/perses/performance_statistics.yaml similarity index 100% rename from examples/online_serving/dashboards/perses/performance_statistics.yaml rename to examples/observability/dashboards/perses/performance_statistics.yaml diff --git a/examples/online_serving/dashboards/perses/query_statistics.yaml b/examples/observability/dashboards/perses/query_statistics.yaml similarity index 100% rename from examples/online_serving/dashboards/perses/query_statistics.yaml rename to examples/observability/dashboards/perses/query_statistics.yaml diff --git a/examples/offline_inference/metrics.py b/examples/observability/metrics/offline.py similarity index 100% rename from examples/offline_inference/metrics.py rename to examples/observability/metrics/offline.py diff --git a/examples/online_serving/opentelemetry/README.md b/examples/observability/opentelemetry/README.md similarity index 100% rename from examples/online_serving/opentelemetry/README.md rename to examples/observability/opentelemetry/README.md diff --git a/examples/online_serving/opentelemetry/dummy_client.py b/examples/observability/opentelemetry/dummy_client.py similarity index 100% rename from examples/online_serving/opentelemetry/dummy_client.py rename to examples/observability/opentelemetry/dummy_client.py diff --git a/examples/online_serving/prometheus_grafana/README.md b/examples/observability/prometheus_grafana/README.md similarity index 100% rename from examples/online_serving/prometheus_grafana/README.md rename to examples/observability/prometheus_grafana/README.md diff --git a/examples/online_serving/prometheus_grafana/docker-compose.yaml b/examples/observability/prometheus_grafana/docker-compose.yaml similarity index 100% rename from examples/online_serving/prometheus_grafana/docker-compose.yaml rename to examples/observability/prometheus_grafana/docker-compose.yaml diff --git a/examples/online_serving/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json similarity index 100% rename from examples/online_serving/prometheus_grafana/grafana.json rename to examples/observability/prometheus_grafana/grafana.json diff --git a/examples/online_serving/prometheus_grafana/prometheus.yaml b/examples/observability/prometheus_grafana/prometheus.yaml similarity index 100% rename from examples/online_serving/prometheus_grafana/prometheus.yaml rename to examples/observability/prometheus_grafana/prometheus.yaml diff --git a/examples/offline_inference/audio_language.py b/examples/offline_inference/audio_language.py index 690aada03aba..c480f1b4145f 100755 --- a/examples/offline_inference/audio_language.py +++ b/examples/offline_inference/audio_language.py @@ -537,9 +537,30 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData: ) +# FireRedLID +def run_fireredlid(question: str, audio_count: int) -> ModelRequestData: + assert audio_count == 1, "FireRedLID only supports single audio input per prompt" + model_name = "PatchyTisa/FireRedLID-vllm" + + prompt = "" + + engine_args = EngineArgs( + model=model_name, + max_model_len=8, + max_num_seqs=5, + limit_mm_per_prompt={"audio": audio_count}, + ) + + return ModelRequestData( + engine_args=engine_args, + prompt=prompt, + ) + + model_example_map = { "audioflamingo3": run_audioflamingo3, "cohere_asr": run_cohere_asr, + "fireredlid": run_fireredlid, "funaudiochat": run_funaudiochat, "gemma3n": run_gemma3n, "glmasr": run_glmasr, diff --git a/examples/offline_inference/encoder_decoder_multimodal.py b/examples/offline_inference/encoder_decoder_multimodal.py index 2f72b7d06705..4fc74e9555f8 100644 --- a/examples/offline_inference/encoder_decoder_multimodal.py +++ b/examples/offline_inference/encoder_decoder_multimodal.py @@ -55,7 +55,91 @@ def run_whisper(): ) +def run_fireredasr2(): + """ + FireRedASR2 – Automatic Speech Recognition model. + + This model uses a Conformer encoder + Qwen2 LLM decoder architecture + for speech-to-text transcription. Audio is passed via the implicit + prompt format with the ``<|AUDIO|>`` placeholder token. + """ + engine_args = EngineArgs( + model="allendou/FireRedASR2-LLM-vllm", + max_model_len=448, + max_num_seqs=16, + limit_mm_per_prompt={"audio": 1}, + ) + + prompt_str = ( + "<|im_start|>user\n<|AUDIO|>请转写音频为文字<|im_end|>\n<|im_start|>assistant\n" + ) + + prompts = [ + { # Implicit prompt with audio + "prompt": prompt_str, + "multi_modal_data": { + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, + }, + }, + { # Another audio sample + "prompt": prompt_str, + "multi_modal_data": { + "audio": AudioAsset("winning_call").audio_and_sample_rate, + }, + }, + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + +def run_fireredlid(): + """ + FireRedLID – Language Identification model. + + This encoder-decoder model identifies the spoken language of an audio + clip. It outputs at most 2 tokens representing the detected language + (e.g. "en", "zh mandarin"). + """ + engine_args = EngineArgs( + model="PatchyTisa/FireRedLID-vllm", + max_model_len=8, + max_num_seqs=16, + limit_mm_per_prompt={"audio": 1}, + ) + + prompts = [ + { # Test explicit encoder/decoder prompt + "encoder_prompt": { + "prompt": "", + "multi_modal_data": { + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, + }, + }, + "decoder_prompt": "", + }, + { # Another audio sample + "encoder_prompt": { + "prompt": "", + "multi_modal_data": { + "audio": AudioAsset("winning_call").audio_and_sample_rate, + }, + }, + "decoder_prompt": "", + }, + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + model_example_map = { + "fireredasr2": run_fireredasr2, + "fireredlid": run_fireredlid, "whisper": run_whisper, } diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index d62f25c2285b..cfeda8804a04 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -394,18 +394,24 @@ def run_eagle2_5(questions: list[str], modality: str) -> ModelRequestData: def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData: model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, max_model_len=4096, max_num_seqs=5, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, trust_remote_code=True, ) + image_placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>" + video_placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>" + if modality == "image": - placeholder = "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>" + placeholder = image_placeholder elif modality == "video": - placeholder = "Video 1:<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>" + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder prompts = [ ( @@ -421,6 +427,49 @@ def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData: ) +# EXAONE-4.5 +def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData: + model_name = "LGAI-EXAONE/EXAONE-4.5-33B" + + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=5, + mm_processor_kwargs={ + "min_pixels": 28 * 28, + "max_pixels": 1280 * 28 * 28, + "fps": 1, + }, + limit_mm_per_prompt=mm_limit, + ) + + image_placeholder = "<|image_pad|>" + video_placeholder = "<|video_pad|>" + + if modality == "image": + placeholder = image_placeholder + elif modality == "video": + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder + + prompts = [ + ( + "<|system|>\nYou are a helpful assistant.<|endofturn|>\n" + f"<|user|>\n{placeholder}" + f"{question}<|endofturn|>\n" + "<|assistant|>\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # Fuyu def run_fuyu(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -529,6 +578,7 @@ def run_glm4v(questions: list[str], modality: str) -> ModelRequestData: def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData: model_name = "zai-org/GLM-4.1V-9B-Thinking" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, max_model_len=4096, @@ -537,14 +587,19 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData: "size": {"shortest_edge": 12544, "longest_edge": 47040000}, "fps": 1, }, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, enforce_eager=True, ) + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + if modality == "image": - placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + placeholder = image_placeholder elif modality == "video": - placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder prompts = [ ( @@ -565,6 +620,7 @@ def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData: def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData: model_name = "zai-org/GLM-4.5V" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, max_model_len=4096, @@ -573,15 +629,20 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData: "size": {"shortest_edge": 12544, "longest_edge": 47040000}, "fps": 1, }, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, enforce_eager=True, tensor_parallel_size=4, ) + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + if modality == "image": - placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + placeholder = image_placeholder elif modality == "video": - placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder prompts = [ ( @@ -602,6 +663,7 @@ def run_glm4_5v(questions: list[str], modality: str) -> ModelRequestData: def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData: model_name = "zai-org/GLM-4.5V-FP8" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, max_model_len=4096, @@ -610,15 +672,20 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData: "size": {"shortest_edge": 12544, "longest_edge": 47040000}, "fps": 1, }, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, enforce_eager=True, tensor_parallel_size=4, ) + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + if modality == "image": - placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + placeholder = image_placeholder elif modality == "video": - placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder prompts = [ ( @@ -639,6 +706,7 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData: def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData: model_name = "zai-org/GLM-OCR" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, max_model_len=4096, @@ -647,14 +715,19 @@ def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData: "size": {"shortest_edge": 12544, "longest_edge": 47040000}, "fps": 1, }, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, enforce_eager=True, ) + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + video_placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + if modality == "image": - placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + placeholder = image_placeholder elif modality == "video": - placeholder = "<|begin_of_video|><|video|><|end_of_video|>" + placeholder = video_placeholder + elif modality == "image+video": + placeholder = image_placeholder + video_placeholder prompts = [ ( @@ -735,11 +808,12 @@ def run_hyperclovax_seed_vision( model_name = "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B" tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, trust_remote_code=True, - max_model_len=8192 if modality == "image" else 16384, - limit_mm_per_prompt={modality: 1}, + max_model_len=16384 if modality in ("video", "image+video") else 8192, + limit_mm_per_prompt=mm_limit, ) messages = list() @@ -791,6 +865,29 @@ def run_hyperclovax_seed_vision( } ] ) + elif modality == "image+video": + messages.append( + [ + { + "role": "user", + "content": [ + { + "type": "image", + "ocr": "", + "lens_keywords": "", + "lens_local_keywords": "", + }, + { + "type": "video", + }, + { + "type": "text", + "text": question, + }, + ], + } + ] + ) else: raise ValueError(f"Unsupported modality: {modality}") @@ -839,19 +936,25 @@ def run_idefics3(questions: list[str], modality: str) -> ModelRequestData: def run_interns1(questions: list[str], modality: str) -> ModelRequestData: model_name = "internlm/Intern-S1-mini" + mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} engine_args = EngineArgs( model=model_name, trust_remote_code=True, max_model_len=8192, max_num_seqs=2, - limit_mm_per_prompt={modality: 1}, + limit_mm_per_prompt=mm_limit, enforce_eager=True, ) + image_placeholder = "" + video_placeholder = "