diff --git a/.buildkite/ci_config_intel.yaml b/.buildkite/ci_config_intel.yaml index a1c0091e0f1..cd6bec005d1 100644 --- a/.buildkite/ci_config_intel.yaml +++ b/.buildkite/ci_config_intel.yaml @@ -2,17 +2,17 @@ name: vllm_intel_ci job_dirs: - ".buildkite/intel_jobs" run_all_patterns: + - ".buildkite/ci_config_intel.yaml" + - ".buildkite/scripts/hardware_ci/run-intel-test.sh" - "docker/Dockerfile" + - "docker/Dockerfile.xpu" - "CMakeLists.txt" - "requirements/common.txt" - "requirements/xpu.txt" - - "requirements/build/cuda.txt" - - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" run_all_exclude_patterns: - - "docker/Dockerfile." - "csrc/cpu/" - "csrc/rocm/" - "cmake/hipify.py" diff --git a/.buildkite/ci_config_rocm.yaml b/.buildkite/ci_config_rocm.yaml index 23f32340071..408ccee58ed 100644 --- a/.buildkite/ci_config_rocm.yaml +++ b/.buildkite/ci_config_rocm.yaml @@ -8,6 +8,7 @@ run_all_patterns: - "docker/docker-bake-rocm.hcl" - ".buildkite/hardware_tests/amd.yaml" - ".buildkite/scripts/ci-bake-rocm.sh" + - ".buildkite/scripts/rocm/" - ".buildkite/scripts/hardware_ci/run-amd-test.py" - ".buildkite/scripts/hardware_ci/run-amd-test.sh" - "CMakeLists.txt" diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index c2510f38aab..d47a7e394a4 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -1,18 +1,45 @@ group: Hardware - AMD Build + +# ROCm image flow: +# 1. Refresh the long-lived ROCm base image only when Dockerfile.rocm_base changes. +# 2. Build ci_base from either the stable base or the freshly refreshed base. +# 3. Build the per-commit ROCm CI image and smoke-test it before GPU jobs run. steps: + - label: "AMD: :docker: refresh ROCm base" + key: refresh-rocm-base-amd + depends_on: [] + device: amd_cpu + no_plugin: true + commands: + - bash .buildkite/scripts/rocm/refresh-base-image.sh + env: + DOCKER_BUILDKIT: "1" + BUILDKIT_PROGRESS: "tty" + TERM: "xterm-256color" + retry: + automatic: + - exit_status: -1 # Agent was lost + limit: 1 + - exit_status: -10 # Agent was lost + limit: 1 + # Ensure ci_base is up-to-date before building the test image. # Compares a content hash of ci_base-affecting files against the remote # image label. If hashes match the build is skipped (< 30 s); if they # differ ci_base is rebuilt and pushed automatically. - label: "AMD: :docker: ensure ci_base" key: ensure-ci-base-amd - depends_on: [] + soft_fail: false + depends_on: + - refresh-rocm-base-amd device: amd_cpu no_plugin: true commands: - - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps + - bash .buildkite/scripts/rocm/build-ci-base.sh env: DOCKER_BUILDKIT: "1" + BUILDKIT_PROGRESS: "tty" + TERM: "xterm-256color" VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl" PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950" REMOTE_VLLM: "1" @@ -26,40 +53,18 @@ steps: - label: "AMD: :docker: build test image and artifacts" key: image-build-amd + soft_fail: false depends_on: - ensure-ci-base-amd device: amd_cpu no_plugin: true commands: - - | - if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then - echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only" - IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts - else - bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel - fi - - | - docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec ' - if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi - if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi - if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi - if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi - command -v python3 - command -v uv - command -v pytest - if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then - echo No ROCm CLI found in image >&2 - exit 1 - fi - python3 - <- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example' + - label: "XPU V1 test" + depends_on: + - image-build-xpu + timeout_in_minutes: 70 + optional: true + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml + - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh v1' + - label: "XPU server test" + depends_on: + - image-build-xpu + timeout_in_minutes: 45 + optional: true + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml + - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server' diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh index 10c03c3e177..52f120534b6 100755 --- a/.buildkite/image_build/image_build.sh +++ b/.buildkite/image_build/image_build.sh @@ -79,12 +79,18 @@ setup_buildx_builder() { docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls } +annotate_image_tags() { + .buildkite/scripts/annotate-image-build.sh \ + "${IMAGE_TAG:-}" "${IMAGE_TAG_LATEST:-}" +} + check_and_skip_if_image_exists() { if [[ -n "${IMAGE_TAG:-}" ]]; then echo "--- :mag: Checking if image exists" if docker manifest inspect "${IMAGE_TAG}" >/dev/null 2>&1; then echo "Image already exists: ${IMAGE_TAG}" echo "Skipping build" + annotate_image_tags exit 0 fi echo "Image not found, proceeding with build" @@ -171,6 +177,18 @@ BRANCH=$4 IMAGE_TAG=$5 IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional +# When TORCH_NIGHTLY=1, build the base CI image against PyTorch nightly so the +# entire existing pipeline runs on nightly torch (CUDA/GPU lane only). Delegate +# to the dedicated nightly build (PYTORCH_NIGHTLY=1, CUDA 13.0) and tag it at the +# normal IMAGE_TAG that every test step already pulls -- no separate image tag, +# no duplicate "vLLM Against PyTorch Nightly" pipeline section. +if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then + echo "--- :warning: TORCH_NIGHTLY=1 -- building base image on PyTorch nightly" + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + exec "${SCRIPT_DIR}/image_build_torch_nightly.sh" \ + "${REGISTRY}" "${REPO}" "${BUILDKITE_COMMIT}" "${BRANCH}" "${IMAGE_TAG}" +fi + # build config TARGET="test-ci" VLLM_BAKE_FILE_PATH="${VLLM_BAKE_FILE_PATH:-docker/docker-bake.hcl}" @@ -254,3 +272,5 @@ echo "--- :docker: Building ${TARGET}" docker --debug buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --progress plain "${TARGET}" echo "--- :white_check_mark: Build complete" + +annotate_image_tags diff --git a/.buildkite/image_build/image_build_arm64.sh b/.buildkite/image_build/image_build_arm64.sh index 5baa55a1965..10cb417ec4d 100755 --- a/.buildkite/image_build/image_build_arm64.sh +++ b/.buildkite/image_build/image_build_arm64.sh @@ -9,29 +9,31 @@ fi REGISTRY=$1 REPO=$2 BUILDKITE_COMMIT=$3 +IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64" # authenticate with AWS ECR 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) ]]; then - echo "Image not found, proceeding with build..." -else +if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then echo "Image found" - exit 0 +else + echo "Image not found, proceeding with build..." + # build for arm64 GPU targets: Grace/GH200 (sm_90), + # Blackwell/Thor (sm_100/sm_103/sm_110), and DGX Spark/GB10 + # (sm_121, family-covered by 12.0 under CUDA 13) + docker build --file docker/Dockerfile \ + --platform linux/arm64 \ + --build-arg max_jobs=16 \ + --build-arg nvcc_threads=4 \ + --build-arg torch_cuda_arch_list="9.0 10.0 11.0 12.0" \ + --build-arg USE_SCCACHE=1 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --tag "$IMAGE" \ + --target test \ + --progress plain . + # push + docker push "$IMAGE" fi -# build (Grace/GH200 is the arm64 GPU target; sm_90) -docker build --file docker/Dockerfile \ - --platform linux/arm64 \ - --build-arg max_jobs=16 \ - --build-arg nvcc_threads=4 \ - --build-arg torch_cuda_arch_list="9.0" \ - --build-arg USE_SCCACHE=1 \ - --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ - --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \ - --target test \ - --progress plain . - -# push -docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 +.buildkite/scripts/annotate-image-build.sh "$IMAGE" diff --git a/.buildkite/image_build/image_build_cpu.sh b/.buildkite/image_build/image_build_cpu.sh index 035f070ab89..3ac1a110f60 100755 --- a/.buildkite/image_build/image_build_cpu.sh +++ b/.buildkite/image_build/image_build_cpu.sh @@ -9,26 +9,26 @@ fi REGISTRY=$1 REPO=$2 BUILDKITE_COMMIT=$3 +IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu" # authenticate with AWS ECR 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 - echo "Image not found, proceeding with build..." -else +if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then echo "Image found" - exit 0 +else + echo "Image not found, proceeding with build..." + # build + docker build --file docker/Dockerfile.cpu \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg VLLM_CPU_X86=true \ + --tag "$IMAGE" \ + --target vllm-test \ + --progress plain . + # push + docker push "$IMAGE" fi -# build -docker build --file docker/Dockerfile.cpu \ - --build-arg max_jobs=16 \ - --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ - --build-arg VLLM_CPU_X86=true \ - --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \ - --target vllm-test \ - --progress plain . - -# push -docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu +.buildkite/scripts/annotate-image-build.sh "$IMAGE" diff --git a/.buildkite/image_build/image_build_cpu_arm64.sh b/.buildkite/image_build/image_build_cpu_arm64.sh index b561e2c2e46..7bb4f10f729 100755 --- a/.buildkite/image_build/image_build_cpu_arm64.sh +++ b/.buildkite/image_build/image_build_cpu_arm64.sh @@ -9,25 +9,25 @@ fi REGISTRY=$1 REPO=$2 BUILDKITE_COMMIT=$3 +IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64-cpu" # authenticate with AWS ECR 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 - echo "Image not found, proceeding with build..." -else +if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then echo "Image found" - exit 0 +else + echo "Image not found, proceeding with build..." + # build + docker build --file docker/Dockerfile.cpu \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --tag "$IMAGE" \ + --target vllm-test \ + --progress plain . + # push + docker push "$IMAGE" fi -# build -docker build --file docker/Dockerfile.cpu \ - --build-arg max_jobs=16 \ - --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ - --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu \ - --target vllm-test \ - --progress plain . - -# push -docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu +.buildkite/scripts/annotate-image-build.sh "$IMAGE" diff --git a/.buildkite/image_build/image_build_hpu.sh b/.buildkite/image_build/image_build_hpu.sh index df900dc6034..35b02c7ec98 100755 --- a/.buildkite/image_build/image_build_hpu.sh +++ b/.buildkite/image_build/image_build_hpu.sh @@ -9,26 +9,26 @@ fi REGISTRY=$1 REPO=$2 BUILDKITE_COMMIT=$3 +IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu" # authenticate with AWS ECR 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"-hpu) ]]; then - echo "Image not found, proceeding with build..." -else +if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then echo "Image found" - exit 0 +else + echo "Image not found, proceeding with build..." + # build + docker build \ + --file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --tag "$IMAGE" \ + --progress plain \ + https://github.com/vllm-project/vllm-gaudi.git + # push + docker push "$IMAGE" fi -# build -docker build \ - --file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \ - --build-arg max_jobs=16 \ - --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ - --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu \ - --progress plain \ - https://github.com/vllm-project/vllm-gaudi.git - -# push -docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu +.buildkite/scripts/annotate-image-build.sh "$IMAGE" diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh index cbd08aa7bd0..d13b3351d3e 100755 --- a/.buildkite/image_build/image_build_torch_nightly.sh +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -40,6 +40,7 @@ docker buildx ls 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" + .buildkite/scripts/annotate-image-build.sh "$IMAGE_TAG" exit 0 fi echo "Image not found, proceeding with build..." @@ -66,3 +67,5 @@ docker buildx build --file docker/Dockerfile \ --progress plain . echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG" + +.buildkite/scripts/annotate-image-build.sh "$IMAGE_TAG" diff --git a/.buildkite/image_build/image_build_xpu.sh b/.buildkite/image_build/image_build_xpu.sh index 45417b7339b..adc544a2331 100755 --- a/.buildkite/image_build/image_build_xpu.sh +++ b/.buildkite/image_build/image_build_xpu.sh @@ -9,26 +9,26 @@ fi REGISTRY=$1 REPO=$2 BUILDKITE_COMMIT=$3 +IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-xpu" # authenticate with AWS ECR 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 # skip build if image already exists -if ! docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu &> /dev/null; then - echo "Image not found, proceeding with build..." -else +if docker manifest inspect "$IMAGE" &> /dev/null; then echo "Image found" - exit 0 +else + echo "Image not found, proceeding with build..." + # build + docker build \ + --file docker/Dockerfile.xpu \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --tag "$IMAGE" \ + --progress plain . + # push + docker push "$IMAGE" fi -# build -docker build \ - --file docker/Dockerfile.xpu \ - --build-arg max_jobs=16 \ - --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ - --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu \ - --progress plain . - -# push -docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu +.buildkite/scripts/annotate-image-build.sh "$IMAGE" diff --git a/.buildkite/intel_jobs/basic_correctness.yaml b/.buildkite/intel_jobs/basic_correctness.yaml index 1b67454d2af..7d9011d6dc8 100644 --- a/.buildkite/intel_jobs/basic_correctness.yaml +++ b/.buildkite/intel_jobs/basic_correctness.yaml @@ -3,8 +3,12 @@ depends_on: - image-build-xpu steps: - label: XPU Sleep Mode - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -19,4 +23,5 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && + pytest -v -s basic_correctness/test_cpu_offload.py && pytest -v -s basic_correctness/test_mem.py::test_end_to_end' diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index c66576d4099..d1dc95b1d40 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -5,6 +5,10 @@ steps: - label: Engine (1 GPU) timeout_in_minutes: 30 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/expert_parallelism_intel.yaml b/.buildkite/intel_jobs/expert_parallelism_intel.yaml index 953e9ddcc55..24dfb07f5f9 100644 --- a/.buildkite/intel_jobs/expert_parallelism_intel.yaml +++ b/.buildkite/intel_jobs/expert_parallelism_intel.yaml @@ -1,11 +1,15 @@ group: Expert Parallelism -depends_on: +depends_on: - image-build-xpu steps: - label: EPLB Algorithm key: eplb-algorithm timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index 66a8db25f02..1407b02055b 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -5,6 +5,10 @@ steps: - label: vLLM IR Tests timeout_in_minutes: 30 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 32a56ef59b3..c1eb1ca63a3 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -5,6 +5,10 @@ steps: - label: LoRA Runtime + Utils timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -34,6 +38,10 @@ steps: - label: LoRA Fused/MoE Kernels timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -54,6 +62,10 @@ steps: - label: LoRA Punica Kernels timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -69,11 +81,17 @@ steps: 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && set -o pipefail && - pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --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]"' + pytest -v -s lora/test_punica_ops.py::test_kernels && + pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size && + pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit' - label: LoRA Punica FP8/XPU Ops - timeout_in_minutes: 45 + timeout_in_minutes: 60 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -94,6 +112,10 @@ steps: - label: LoRA Models timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -108,15 +130,19 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - (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_transformers_model.py && pytest -v -s lora/test_chatglm3_tp.py && + pytest -v -s lora/test_llama_tp.py::test_llama_lora && pytest -s -v lora/test_minicpmv_tp.py' - label: LoRA Multimodal timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 20a9f5bd173..047b00c49c7 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -3,8 +3,12 @@ depends_on: - image-build-xpu steps: - label: V1 Core + KV + Metrics - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -29,8 +33,12 @@ steps: pytest -v -s v1/executor' - label: V1 Sample + Logits - timeout_in_minutes: 30 + timeout_in_minutes: 90 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -57,17 +65,46 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && + 'pip install lm_eval[api]>=0.4.12 && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py && pytest -v -s v1/test_oracle.py && pytest -v -s v1/test_request.py && pytest -v -s v1/test_outputs.py && - pytest -v -s v1/sample/test_topk_topp_sampler.py' + pytest -v -s v1/sample' + +- label: Basic Models Tests (Initialization) + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + 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/models/test_initialization.py + - tests/models/registry.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_XPU_FUSED_MOE_USE_REF=1 && + cd tests && + pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' - label: XPU CPU Offload timeout_in_minutes: 60 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -88,10 +125,39 @@ steps: pytest -v -s v1/kv_offload && pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' +- label: NixlConnector PD accuracy (2 GPUs) + timeout_in_minutes: 60 + num_devices: 2 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/xpu.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' + - label: Regression key: regression - timeout_in_minutes: 30 + timeout_in_minutes: 50 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -114,15 +180,19 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install modelscope && + 'pip install modelscope\<1.38 && cd tests && pytest -v -s test_regression.py' - label: Metrics, Tracing (2 GPUs) key: metrics-tracing-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 45 num_devices: 2 device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -152,8 +222,12 @@ steps: - label: Async Engine, Inputs, Utils, Worker key: async-engine-inputs-utils-worker - timeout_in_minutes: 30 + timeout_in_minutes: 55 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml new file mode 100644 index 00000000000..0311b5dffb7 --- /dev/null +++ b/.buildkite/intel_jobs/model_runner_v2_intel.yaml @@ -0,0 +1,62 @@ +group: Model Runner V2 Intel +depends_on: + - image-build-xpu +steps: +- label: Model Runner V2 Core Tests (Intel) + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + 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/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - vllm/v1/core/sched/ + - vllm/v1/attention/ + - tests/v1/engine/test_llm_engine.py + - tests/v1/e2e/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd tests && + pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" && + ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" && + pytest -v -s v1/e2e/general/test_min_tokens.py' + +- label: Model Runner V2 Examples (Intel) + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + 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/worker/gpu/ + - vllm/v1/core/sched/ + - vllm/v1/worker/gpu_worker.py + - examples/basic/offline_inference/ + - examples/generate/multimodal/ + - examples/features/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd examples && + python3 basic/offline_inference/chat.py && + python3 basic/offline_inference/generate.py --model facebook/opt-125m && + python3 generate/multimodal/vision_language_offline.py --seed 0 && + python3 features/automatic_prefix_caching/prefix_caching_offline.py' diff --git a/.buildkite/intel_jobs/models_distributed_intel.yaml b/.buildkite/intel_jobs/models_distributed_intel.yaml new file mode 100644 index 00000000000..5f2247eeaac --- /dev/null +++ b/.buildkite/intel_jobs/models_distributed_intel.yaml @@ -0,0 +1,27 @@ +group: Models - Distributed +depends_on: + - image-build-xpu +steps: +- label: Distributed Model Tests (2 GPUs) + key: distributed-model-tests-2-gpus + timeout_in_minutes: 65 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/model_executor/model_loader/sharded_state_loader.py + - vllm/model_executor/models/ + - tests/model_executor/model_loader/test_sharded_state_loader.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m "not slow_test"' diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index cf5b51c4b89..6d04662fa5a 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -1,11 +1,15 @@ group: Models - Multimodal -depends_on: +depends_on: - image-build-xpu steps: - label: "Multi-Modal Models (Standard) 1: qwen2" key: multi-modal-models-standard-1-qwen2 - timeout_in_minutes: 45 + timeout_in_minutes: 70 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -18,15 +22,19 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git && + 'pip install av && cd tests && 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" key: multi-modal-models-standard-2-qwen3-gemma - timeout_in_minutes: 45 + timeout_in_minutes: 70 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -39,14 +47,17 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git && - cd tests && + 'cd tests && pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model' - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" key: multi-modal-models-standard-3-llava-qwen2-vl - timeout_in_minutes: 45 + timeout_in_minutes: 65 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -59,8 +70,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git && - cd tests && + 'cd tests && 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' @@ -68,6 +78,10 @@ steps: key: multi-modal-models-standard-4-other-whisper timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -80,14 +94,18 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git && + 'pip install av && cd tests && 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' - label: Multi-Modal Processor # 44min key: multi-modal-processor - timeout_in_minutes: 45 + timeout_in_minutes: 60 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -101,11 +119,9 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git && + 'pip install av matplotlib ftfy && pip install open-clip-torch --no-deps && cd tests && pytest -v -s models/multimodal/processing/test_tensor_schema.py - --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]" - --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' parallelism: 4 diff --git a/.buildkite/intel_jobs/quantization.yaml b/.buildkite/intel_jobs/quantization.yaml new file mode 100644 index 00000000000..a5c11da9dd2 --- /dev/null +++ b/.buildkite/intel_jobs/quantization.yaml @@ -0,0 +1,28 @@ +group: Quantization +depends_on: + - image-build-xpu +steps: +- label: Quantization + key: quantization + timeout_in_minutes: 30 + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + no_plugin: true + working_dir: "." + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/quantization + commands: + # - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]"' + diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 4bf14f7064b..ec58512441d 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -17,8 +17,12 @@ steps: - label: "XPU example Test" depends_on: - image-build-xpu - timeout_in_minutes: 30 + timeout_in_minutes: 50 device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -38,17 +42,46 @@ steps: 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 nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 && 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 TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager && 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 && python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 && VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192 ' + - label: "XPU W8A8 FP8 Linear Examples" + depends_on: + - image-build-xpu + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 && + python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 && + python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096 + ' - label: "XPU V1 test" depends_on: - image-build-xpu - timeout_in_minutes: 30 + timeout_in_minutes: 70 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -66,13 +99,17 @@ steps: 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/test_serial_utils.py && - pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py && + pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.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 --ignore=v1/kv_connector/unit/test_offloading_connector.py' - label: "XPU server test" depends_on: - image-build-xpu - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -87,3 +124,47 @@ steps: cd tests && pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py && pytest -v -s benchmarks/test_serve_cli.py' + - label: "XPU quantization test" + depends_on: + - image-build-xpu + timeout_in_minutes: 30 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - vllm/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s quantization/test_auto_round.py' + - label: "XPU compressed tensors FP8 test" + depends_on: + - image-build-xpu + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/quantization/test_compressed_tensors.py + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8' \ No newline at end of file diff --git a/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml index a87328fcdcc..164733cca6f 100644 --- a/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml +++ b/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml @@ -6,9 +6,7 @@ tasks: value: 0.7142 - name: "exact_match,flexible-extract" value: 0.4579 -env_vars: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +moe_backend: "flashinfer_cutlass" limit: 1319 num_fewshot: 5 max_model_len: 262144 diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index d34e603b9e2..dd2fd5f05b4 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -68,6 +68,10 @@ def launch_lm_eval(eval_config, tp_size): if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]: model_args += "attention_backend=TRITON_ATTN" + moe_backend = eval_config.get("moe_backend", None) + if moe_backend is not None: + model_args += f"moe_backend={moe_backend}," + env_vars = eval_config.get("env_vars", None) with scoped_env_vars(env_vars): results = lm_eval.simple_evaluate( diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index a34f534e54d..1d3a1d4a534 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -1,12 +1,25 @@ # 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 +# These requested arches are filtered by CMake's CUDA_SUPPORTED_ARCHS before +# per-kernel arch selection. Do not add +PTX here: top-level +PTX is stripped +# during that filtering, so kernels that need PTX must request it locally. 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" + # for CUDA >=13, sm_100+ targets have family specifiers (see CMakeLists.txt) + # so targets like 10.3 and 12.1 are automatically supported with this list + CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0" + # aarch64-only targets: Orin (8.7), Thor (11.0, CUDA 13+) + CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0" + + # for CUDA <13, we need to specify all needed targets + # some targets (10.3, 12.1) are skipped to limit the wheel size (< 500MB) + # please use CUDA 13 wheels or compile yourself on these new devices 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" + + # pre-built mooncake wheels + # the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04 + # so we use different wheels for the time being MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl" MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl" MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl" @@ -337,6 +350,25 @@ steps: - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' + - label: ":docker: Build release image - x86_64 - XPU" + depends_on: ~ + id: build-xpu-release-image + 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 \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh xpu) \ + --build-arg GIT_REPO_CHECK=1 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile.xpu . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-xpu" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-xpu"' + env: + DOCKER_BUILDKIT: "1" + - block: "Build release image for x86_64 CPU" key: block-cpu-release-image-build depends_on: ~ @@ -432,6 +464,16 @@ steps: - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"' + - label: "Create manifest - XPU" + depends_on: + - build-xpu-release-image + id: create-manifest-xpu + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/scripts/xpu/create-xpu-ecr-manifest.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: XPU" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-xpu"' + - label: "Publish nightly multi-arch image to DockerHub" depends_on: - create-multi-arch-manifest @@ -806,6 +848,23 @@ steps: DOCKER_BUILDKIT: "1" S3_BUCKET: "vllm-wheels" + - label: "Publish nightly XPU image to DockerHub" + depends_on: + - create-manifest-xpu + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" + plugins: + - docker-login#v3.0.0: + username: vllmbot + password-env: DOCKERHUB_TOKEN + env: + DOCKER_BUILDKIT: "1" + DOCKERHUB_USERNAME: "vllmbot" + - label: "Publish nightly ROCm image to DockerHub" depends_on: - build-rocm-release-image @@ -836,6 +895,7 @@ steps: - create-multi-arch-manifest-cuda-12-9 - create-multi-arch-manifest-ubuntu2404 - create-multi-arch-manifest-cuda-12-9-ubuntu2404 + - create-manifest-xpu - build-rocm-release-image - input-release-version # Wait for CPU builds if their block steps were unblocked, so publish @@ -846,7 +906,6 @@ steps: allow_failure: true - step: build-cpu-release-image-arm64 allow_failure: true - if: build.env("NIGHTLY") != "1" - label: "Publish release images to DockerHub" depends_on: diff --git a/.buildkite/scripts/annotate-image-build.sh b/.buildkite/scripts/annotate-image-build.sh new file mode 100755 index 00000000000..174d88bf674 --- /dev/null +++ b/.buildkite/scripts/annotate-image-build.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Append the Docker image tag(s) an image-build step pushed to a Buildkite +# annotation, so the built image tags show up on the build page instead of +# being buried in the job logs. +# +# Usage: annotate-image-build.sh [ ...] +set -euo pipefail + +# buildkite-agent only exists on Buildkite agents; no-op elsewhere so the +# image build scripts stay runnable locally. +if ! command -v buildkite-agent >/dev/null 2>&1; then + echo "buildkite-agent not found; skipping image tag annotation" + exit 0 +fi + +label="${BUILDKITE_LABEL:-Image build}" +content="" +for image in "$@"; do + [[ -n "$image" ]] || continue + content+="- **${label}**: \`${image}\`"$'\n' +done + +if [[ -z "$content" ]]; then + echo "No image tags provided; nothing to annotate" + exit 0 +fi + +# Best-effort: a flaky annotation must never fail an otherwise successful +# (and expensive) image build. +if ! printf '%s' "$content" | \ + buildkite-agent annotate --append --style 'info' --context 'docker-images'; then + echo "warning: failed to annotate build with image tags" +fi diff --git a/.buildkite/scripts/check-ray-compatibility.sh b/.buildkite/scripts/check-ray-compatibility.sh index b056d4403db..ec42834802d 100644 --- a/.buildkite/scripts/check-ray-compatibility.sh +++ b/.buildkite/scripts/check-ray-compatibility.sh @@ -29,7 +29,11 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then TORCH_INDEX_URL="" fi else - TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130" + if [ "${TORCH_NIGHTLY:-0}" = "1" ]; then + TORCH_INDEX_URL="https://download.pytorch.org/whl/nightly/cu130" + else + TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130" + fi fi echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}" diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 1289939180d..22391291b44 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -15,9 +15,10 @@ set -euo pipefail DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" -DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base tools/install_torchcodec_rocm.sh tests/vllm_test_utils" +DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_METADATA_VERSION="1" IMAGE_EXISTED_BEFORE_BUILD=0 TARGET="" @@ -392,6 +393,16 @@ should_upload_wheel_artifacts() { || "${TARGET}" == *"artifact"* ]] } +set_buildkite_metadata() { + local key="$1" + local value="$2" + + [[ -n "${value}" ]] || return 0 + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data set "${key}" "${value}" || true + fi +} + get_remote_image_label() { local image_ref="$1" local label_key="$2" @@ -525,6 +536,22 @@ get_remote_image_label_with_retry() { return 0 } +remote_ci_base_metadata_is_current() { + local image_ref="$1" + local metadata_version="" + + metadata_version=$(get_remote_image_label "${image_ref}" "vllm.ci_base.metadata_version") + [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]] +} + +remote_ci_base_metadata_is_current_with_retry() { + local image_ref="$1" + local metadata_version="" + + metadata_version=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.metadata_version") + [[ "${metadata_version}" == "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" ]] +} + remote_image_exists() { local image_ref="$1" docker manifest inspect "${image_ref}" >/dev/null 2>&1 @@ -581,6 +608,7 @@ init_config() { CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" export PYTORCH_ROCM_ARCH @@ -635,6 +663,10 @@ load_ci_hcl() { echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}" } +init_bake_files() { + BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") +} + compute_ci_base_hash_if_needed() { if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then return 0 @@ -676,12 +708,14 @@ configure_ci_base_image_refs() { fi content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}") + CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}" if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then commit_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${BUILDKITE_COMMIT}") - CI_BASE_IMAGE_TAG_COMMIT="${commit_tag}" - export CI_BASE_IMAGE_TAG_COMMIT fi + CI_BASE_IMAGE_TAG_COMMIT_REF="${commit_tag}" + # *_REF is the logical tag recorded in metadata. *_EXTRA is only passed to + # bake when that tag is not already the primary tag, avoiding duplicates. if should_push_stable_ci_base_tag; then primary_tag="${content_tag}" CI_BASE_IMAGE_TAG_STABLE="${stable_tag}" @@ -691,19 +725,35 @@ configure_ci_base_image_refs() { fi CI_BASE_IMAGE_TAG="${primary_tag}" if [[ "${primary_tag}" == "${content_tag}" ]]; then - CI_BASE_IMAGE_TAG_CONTENT="" + CI_BASE_IMAGE_TAG_CONTENT_EXTRA="" else - CI_BASE_IMAGE_TAG_CONTENT="${content_tag}" + CI_BASE_IMAGE_TAG_CONTENT_EXTRA="${content_tag}" fi - export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE + if [[ -n "${commit_tag}" && "${commit_tag}" != "${primary_tag}" ]]; then + CI_BASE_IMAGE_TAG_COMMIT_EXTRA="${commit_tag}" + else + CI_BASE_IMAGE_TAG_COMMIT_EXTRA="" + fi + export CI_BASE_IMAGE_TAG + export CI_BASE_IMAGE_TAG_COMMIT_EXTRA + export CI_BASE_IMAGE_TAG_CONTENT_EXTRA + export CI_BASE_IMAGE_TAG_CONTENT_REF + export CI_BASE_IMAGE_TAG_COMMIT_REF + export CI_BASE_IMAGE_TAG_STABLE if is_ci_base_target; then IMAGE_TAG="${primary_tag}" + CI_BASE_IMAGE="${primary_tag}" + export CI_BASE_IMAGE export IMAGE_TAG echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}" - if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then - echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}" + if [[ -n "${commit_tag}" ]]; then + if [[ "${commit_tag}" == "${primary_tag}" ]]; then + echo "ci_base commit image tag: ${commit_tag} (primary)" + else + echo "ci_base commit image tag: ${commit_tag}" + fi fi echo "ci_base content image tag: ${content_tag}" if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then @@ -712,6 +762,10 @@ configure_ci_base_image_refs() { echo "ci_base stable alias will not be pushed for this build" echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}" fi + set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}" + set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}" + set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-}" + set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}" return 0 fi @@ -728,8 +782,8 @@ ci_base_candidate_refs() { printf '%s\n' \ "${IMAGE_TAG:-}" \ "${CI_BASE_IMAGE_TAG:-}" \ - "${CI_BASE_IMAGE_TAG_COMMIT:-}" \ - "${CI_BASE_IMAGE_TAG_CONTENT:-}" \ + "${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}" \ + "${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}" \ "${CI_BASE_IMAGE_TAG_STABLE:-}" \ | awk 'NF && !seen[$0]++' } @@ -743,6 +797,10 @@ find_matching_ci_base_ref() { remote_image_exists "${candidate}" || continue candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash") if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + if ! remote_ci_base_metadata_is_current "${candidate}"; then + echo "Found matching ci_base content hash but stale metadata: ${candidate}" >&2 + continue + fi printf '%s\n' "${candidate}" return 0 fi @@ -817,6 +875,10 @@ maybe_skip_existing_image() { if [[ -n "${remote_hash}" ]]; then echo "Remote ci_base content hash: ${remote_hash:0:16}..." if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + if ! remote_ci_base_metadata_is_current "${IMAGE_TAG}"; then + echo "Content hashes match but ci_base metadata is stale; rebuilding to refresh metadata" + return 0 + fi if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then echo "ci_base tag refresh failed; rebuilding to push expected tags" return 0 @@ -998,12 +1060,104 @@ prepare_git_cache_metadata() { fi } +ci_base_metadata_pairs() { + local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" + local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" + local content_files_hash="" + local base_image="" + local base_image_digest="" + local git_branch="" + local -a content_paths=() + local -a content_args=() + + read -r -a content_paths <<< "${content_files}" + if [[ ${#content_paths[@]} -gt 0 ]]; then + content_files_hash=$(compute_content_hash "${content_paths[@]}") + fi + mapfile -t content_args < <( + get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}" + ) + + base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE") + if [[ -n "${base_image}" ]]; then + base_image_digest=$(resolve_image_digest "${base_image}") + fi + git_branch="${BUILDKITE_BRANCH:-${VLLM_BRANCH:-}}" + + metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" + metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}" + metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}" + metadata_pair "vllm.ci_base.content_files" "${content_files}" + metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")" + metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}" + metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}" + metadata_pair "vllm.ci_base.image.primary" "${CI_BASE_IMAGE_TAG:-}" + metadata_pair "vllm.ci_base.image.content" "${CI_BASE_IMAGE_TAG_CONTENT_REF:-${CI_BASE_IMAGE_TAG_CONTENT_EXTRA:-}}" + metadata_pair "vllm.ci_base.image.commit" "${CI_BASE_IMAGE_TAG_COMMIT_REF:-${CI_BASE_IMAGE_TAG_COMMIT_EXTRA:-}}" + metadata_pair "vllm.ci_base.image.stable" "${CI_BASE_IMAGE_TAG_STABLE:-}" + metadata_pair "vllm.ci_base.git_commit" "${BUILDKITE_COMMIT:-}" + metadata_pair "vllm.ci_base.git_branch" "${git_branch}" + metadata_pair "vllm.ci_base.vllm_branch" "${VLLM_BRANCH:-}" + metadata_pair "vllm.ci_base.stable_branch" "${CI_BASE_STABLE_BRANCH:-main}" + + metadata_pair "vllm.rocm.base_image" "${base_image}" + metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}" + metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}" + metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")" + metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")" + metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")" + metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")" + metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}" + metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")" + metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}" + metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")" + metadata_pair "vllm.rocm.rocshmem_commit" "${ROCSHMEM_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_BRANCH")}" + metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")" + metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" + metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" + metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" + metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}" + metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}" + metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" + + metadata_pair "vllm.buildkite.build_number" "${BUILDKITE_BUILD_NUMBER:-}" + metadata_pair "vllm.buildkite.build_id" "${BUILDKITE_BUILD_ID:-}" +} + +write_ci_base_metadata_annotations() { + local metadata="$1" + local key="" + local value="" + local annotation="" + + [[ -n "${metadata}" ]] || return 0 + while IFS=$'\t' read -r key value; do + [[ -n "${key}" && -n "${value}" ]] || continue + annotation="manifest:${key}=${value}" + printf ' "%s",\n' "$(hcl_escape_string "${annotation}")" + done <<< "${metadata}" +} + +write_ci_base_metadata_labels() { + local metadata="$1" + local key="" + local value="" + + [[ -n "${metadata}" ]] || return 0 + while IFS=$'\t' read -r key value; do + [[ -n "${key}" && -n "${value}" ]] || continue + printf ' "%s" = "%s"\n' \ + "$(hcl_escape_string "${key}")" \ + "$(hcl_escape_string "${value}")" + done <<< "${metadata}" +} + write_ci_base_label_override() { local target_name="" + local metadata="" local -a ci_base_targets=() - BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") - if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then return 0 fi @@ -1019,16 +1173,23 @@ write_ci_base_label_override() { return 0 fi + metadata=$(ci_base_metadata_pairs) + : > "${CI_BASE_LABEL_OVERRIDE_PATH}" for target_name in "${ci_base_targets[@]}"; do cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" </dev/null 2>&1; then + echo "${cache_name} content cache exists; not re-exporting ${cache_ref}" + return 1 + fi + echo "${cache_name} content cache missing; will export ${cache_ref}" + return 0 + ;; + *) + echo "Error: ROCM_CONTENT_CACHE_EXPORT_MODE must be one of: missing, always, never" + exit 1 + ;; + esac +} + write_rocm_cache_override() { local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" + local content_cache_export_mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}" local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}" + local rust_cache_to_mode="${ROCM_RUST_CACHE_TO_MODE:-max}" local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}" - local -a content_cache_from=() + local -a csrc_content_cache_from=() + local -a rust_content_cache_from=() + local -a combined_content_cache_from=() local -a csrc_cache_to=() + local -a rust_cache_to=() local -a rocm_cache_to=() local -a export_wheel_cache_to=() + local export_csrc_cache=1 + local export_rust_cache=1 - if ! uses_rocm_csrc_cache; then + if ! uses_rocm_csrc_cache && ! uses_rocm_rust_cache; then return 0 fi + validate_content_cache_export_mode \ + "${content_cache_export_mode}" \ + "ROCM_CONTENT_CACHE_EXPORT_MODE" validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE" + validate_cache_export_mode "${rust_cache_to_mode}" "ROCM_RUST_CACHE_TO_MODE" validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE" + echo "ROCm content cache export mode: ${content_cache_export_mode}" echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}" + echo "ROCm Rust cache export mode: ${rust_cache_to_mode}" echo "ROCm final image cache export mode: ${rocm_cache_to_mode}" if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then - content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}") - csrc_cache_to+=( - "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true" - ) + csrc_content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}") + if should_export_content_cache_ref "${ROCM_CSRC_CONTENT_CACHE_REF}" "ROCm csrc"; then + csrc_cache_to+=( + "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true" + ) + else + export_csrc_cache=0 + fi fi + if [[ -n "${ROCM_RUST_CONTENT_CACHE_REF:-}" ]]; then + rust_content_cache_from+=("type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}") + if should_export_content_cache_ref "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust"; then + rust_cache_to+=( + "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=${rust_cache_to_mode},ignore-error=true" + ) + else + export_rust_cache=0 + fi + fi + + combined_content_cache_from=("${csrc_content_cache_from[@]}" "${rust_content_cache_from[@]}") + # Docker Hub cache exports are best-effort. A cache-only target failure can # otherwise cancel the sibling image target before its manifest is pushed. if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then - csrc_cache_to+=( - "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true" - ) + if [[ ${export_csrc_cache} -eq 1 ]]; then + csrc_cache_to+=( + "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true" + ) + fi + if [[ ${export_rust_cache} -eq 1 ]]; then + rust_cache_to+=( + "type=registry,ref=${cache_repo}:rust-rocm-${BUILDKITE_COMMIT},mode=${rust_cache_to_mode},ignore-error=true" + ) + fi rocm_cache_to+=( "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true" ) fi if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then - csrc_cache_to+=( - "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true" - ) + if [[ ${export_csrc_cache} -eq 1 ]]; then + csrc_cache_to+=( + "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true" + ) + fi + if [[ ${export_rust_cache} -eq 1 ]]; then + rust_cache_to+=( + "type=registry,ref=${cache_repo}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rust_cache_to_mode},ignore-error=true" + ) + fi rocm_cache_to+=( "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true" ) @@ -1249,7 +1572,7 @@ target "csrc-rocm-ci" { cache-from = concat( get_cache_from_rocm_csrc(), EOF - write_hcl_string_list " " "${content_cache_from[@]}" + write_hcl_string_list " " "${csrc_content_cache_from[@]}" cat < BEL) +sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE" + # Strip colorization sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE" diff --git a/.buildkite/scripts/ci-fetch-log.sh b/.buildkite/scripts/ci-fetch-log.sh index 3f99bc50a57..4830135a112 100755 --- a/.buildkite/scripts/ci-fetch-log.sh +++ b/.buildkite/scripts/ci-fetch-log.sh @@ -1,74 +1,178 @@ #!/bin/bash -# Usage: ./ci-fetch-log.sh [output_file] -# ./ci-fetch-log.sh [output_file] +# Fetch vLLM Buildkite CI logs (public; no login required). # -# Downloads the raw log for a Buildkite job from the public, unauthenticated -# /organizations//pipelines//builds//jobs//download -# endpoint, then strips ANSI/timestamps via ci-clean-log.sh. +# Usage: +# ci-fetch-log.sh [--soft|--all] --pr [] failed jobs in the PR's latest +# build (current branch if omitted) +# ci-fetch-log.sh [--soft|--all] failed jobs in that build +# ci-fetch-log.sh [output] one job; both # and +# ?sid= URL forms work +# ci-fetch-log.sh [output] # -# Find and via: -# gh pr checks --repo vllm-project/vllm -# Each failing row's URL is .../builds/#. -# -# Default output path: ci--.log (e.g. -# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's -# first 8 chars, so the second segment is needed for uniqueness when -# fetching multiple jobs in parallel. The script refuses to overwrite an -# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1 -# to override. +# --soft also fetches soft-failed jobs; --all fetches every finished job. +# Saves each log as ci--.log (ANSI/timestamps stripped) and +# prints "\t" per job. [output] is single-job only; "-" +# streams to stdout. Existing files are kept; CI_FETCH_LOG_FORCE=1 refetches. set -euo pipefail ORG="vllm" PIPELINE="ci" +UA="vllm-ci-fetch-log" +UUID_RE='[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}' usage() { - echo "Usage: $0 [output_file]" - echo " $0 [output_file]" + sed -n '2,15p' "$0" | sed 's/^# \{0,1\}//' exit 1 } -if [ $# -lt 1 ]; then usage; fi +die() { + echo "$1" >&2 + exit 1 +} -if [[ "$1" == https://* ]]; then +BUILD="" JOB="" SID="" OUT="" +SCOPE="failed" + +while :; do + case "${1:-}" in + --soft) SCOPE="soft" ;; + --all) SCOPE="all" ;; + *) break ;; + esac + shift +done + +case "${1:-}" in +--pr) + PR="${2:-}" + # gh pr checks exits non-zero when checks are failing; that is the + # expected case here. + URL=$(gh pr checks ${PR:+"$PR"} --repo vllm-project/vllm 2>/dev/null | + grep -oE "https://buildkite.com/${ORG}/${PIPELINE}/builds/[0-9]+" | + sort -t/ -k7 -n | tail -1 || true) + [ -n "$URL" ] || die "No Buildkite build found via: gh pr checks ${PR:-}" + BUILD="${URL##*/}" + ;; +https://*) BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p') - JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1) + JOB=$(echo "$1" | grep -oE "#${UUID_RE}" | head -n 1 | cut -c2- || true) + SID=$(echo "$1" | grep -oE "[?&]sid=${UUID_RE}" | head -n 1 | sed 's/.*sid=//' || true) OUT="${2:-}" -else - if [ $# -lt 2 ]; then usage; fi + [ -n "$BUILD" ] || die "Could not parse build number from: $1" + ;; +[0-9]*) + [ $# -ge 2 ] || usage BUILD="$1" JOB="$2" OUT="${3:-}" -fi - -if [ -z "$BUILD" ] || [ -z "$JOB" ]; then - echo "Could not parse build number or job UUID from: $1" >&2 + ;; +*) usage -fi - -# Jobs in the same build share the UUID's first segment, so include the -# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames -# unique when fetching multiple jobs from one build in parallel. -if [ -z "$OUT" ]; then - OUT="ci-${BUILD}-${JOB:0:13}.log" -fi - -if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then - echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2 - exit 1 -fi + ;; +esac COOKIES=$(mktemp) -trap 'rm -f "$COOKIES"' EXIT +JOBS_TSV=$(mktemp) +trap 'rm -f "$COOKIES" "$JOBS_TSV"' EXIT -# Buildkite issues a session cookie on first hit; subsequent /download needs it. -curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \ +# Buildkite issues a session cookie on first hit; later requests need it. +curl -fsSL -c "$COOKIES" -A "$UA" \ "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null -curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \ - "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \ - -o "$OUT" +# The build's job list (id, step uuid, state, name) is served as JSON from +# the user-facing /data/jobs endpoint. Flatten it to TSV for easy filtering: +# job_id step_uuid failed soft_failed finished slug name +curl -fsSL -b "$COOKIES" -A "$UA" \ + "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}/data/jobs" | + python3 -c ' +import json, re, sys -bash "$(dirname "$0")/ci-clean-log.sh" "$OUT" +data = json.load(sys.stdin) +if data.get("has_next_page"): + print("warning: job list is paginated; some jobs not shown", file=sys.stderr) +for r in data["records"]: + if r.get("type") != "script": + continue + name = (r.get("name") or "").replace("\t", " ").replace("\n", " ") + slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:60] + print("\t".join([ + r["id"], + r.get("step_uuid") or "", + str(r.get("passed") is False), + str(bool(r.get("soft_failed"))), + str(bool(r.get("finished_at"))), + slug, + name, + ])) +' >"$JOBS_TSV" || die "Could not list jobs for build ${BUILD}" -echo "$OUT" +if [ -n "$SID" ] && [ -z "$JOB" ]; then + # The ?sid= in builds//list URLs is the *step* uuid, not the job uuid. + JOB=$(awk -F'\t' -v s="$SID" '$1 == s || $2 == s {print $1; exit}' "$JOBS_TSV") + [ -n "$JOB" ] || die "No job matching sid=${SID} in build ${BUILD}" +fi + +fetch_job() { # + curl -fsSL -b "$COOKIES" -A "$UA" \ + "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/$1/download" \ + -o "$2" + bash "$(dirname "$0")/ci-clean-log.sh" "$2" +} + +if [ -n "$JOB" ]; then + # Single-job mode. + NAME=$(awk -F'\t' -v j="$JOB" '$1 == j {print $7; exit}' "$JOBS_TSV") + SLUG=$(awk -F'\t' -v j="$JOB" '$1 == j {print $6; exit}' "$JOBS_TSV") + [ -n "$OUT" ] || OUT="ci-${BUILD}-${SLUG:-${JOB:0:13}}.log" + if [ "$OUT" = "-" ]; then + TMP=$(mktemp) + fetch_job "$JOB" "$TMP" + cat "$TMP" + rm -f "$TMP" + exit 0 + fi + if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then + die "Refusing to overwrite existing ${OUT} (set CI_FETCH_LOG_FORCE=1 or pass an output path)." + fi + fetch_job "$JOB" "$OUT" + printf '%s\t%s\n' "$OUT" "${NAME:-$JOB}" + exit 0 +fi + +# Build-wide mode: fetch finished jobs matching $SCOPE. +[ -z "$OUT" ] || die "[output_file] is only valid when fetching a single job." + +case "$SCOPE" in +failed) FILTER='$3 == "True" && $4 == "False" && $5 == "True"' ;; +soft) FILTER='$3 == "True" && $5 == "True"' ;; +all) FILTER='$5 == "True"' ;; +esac + +if [ "$SCOPE" = "failed" ]; then + SOFT=$(awk -F'\t' '$3 == "True" && $4 == "True"' "$JOBS_TSV" | wc -l) + [ "$SOFT" -eq 0 ] || echo "Skipping ${SOFT} soft-failed job(s); use --soft to include them." >&2 +fi + +FOUND=0 +EMITTED=" " +while IFS=$'\t' read -r job_id _ _ _ _ slug name; do + FOUND=$((FOUND + 1)) + out="ci-${BUILD}-${slug:-${job_id:0:13}}.log" + # Retries share a name with the original job; disambiguate by uuid. + case "$EMITTED" in + *" $out "*) out="ci-${BUILD}-${slug:-job}-${job_id:0:13}.log" ;; + esac + EMITTED="${EMITTED}${out} " + if [ -e "$out" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then + echo "Keeping existing ${out} (set CI_FETCH_LOG_FORCE=1 to refetch)." >&2 + elif ! fetch_job "$job_id" "$out"; then + echo "Failed to download log for job ${job_id} (${name})." >&2 + continue + fi + printf '%s\t%s\n' "$out" "$name" +done < <(awk -F'\t' "$FILTER" "$JOBS_TSV") + +if [ "$FOUND" -eq 0 ]; then + echo "No matching jobs in build ${BUILD} (scope: ${SCOPE})." >&2 +fi diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 5c994e25d0c..fee9ab04f4b 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -28,6 +28,17 @@ ############################################################################### set -o pipefail +: "${BUILDKIT_PROGRESS:=plain}" +: "${TERM:=xterm-256color}" +: "${FORCE_COLOR:=1}" +: "${CLICOLOR_FORCE:=1}" +: "${PY_COLORS:=1}" +: "${ROCM_DOCKER_TTY:=1}" +if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then + PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes" +fi +export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY + # Export Python path for commands that run directly on the host. Containerized # tests set this to /vllm-workspace below so spawned Python processes do not # depend on their current working directory. @@ -149,6 +160,7 @@ EOF echo "--- Building local ROCm test image" docker build \ --pull=false \ + --progress "${BUILDKIT_PROGRESS}" \ --build-arg "BASE_IMAGE=${base_image}" \ -t "${artifact_image}" \ "${context_dir}" || return 1 @@ -367,6 +379,20 @@ remove_docker_container() { } trap remove_docker_container EXIT +# python_only_compile.sh runs `python setup.py develop` and needs the full repo tree +# under /vllm-workspace (Dockerfile.rocm test stage: mkdir src && mv vllm). +# The ROCm wheel artifact tarball only ships a thin tree (tests, etc.), so +# artifact images cannot satisfy that test — use the full rocm/vllm-ci image. +_cmd_probe="${VLLM_TEST_COMMANDS:-}" +if [[ -z "${_cmd_probe}" ]]; then + _cmd_probe="$*" +fi +if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" == "1" && "${_cmd_probe}" == *python_only_compile.sh* ]]; then + echo "INFO: disabling VLLM_CI_USE_ARTIFACTS for python_only_compile (requires full /vllm-workspace tree)" + export VLLM_CI_USE_ARTIFACTS=0 +fi +unset -v _cmd_probe + if ! prepare_artifact_image; then echo "Using full ROCm CI image: ${image_name}" docker pull "${image_name}" || exit 1 @@ -426,6 +452,26 @@ fi echo "Final commands: $commands" +# The ROCm test image often ships /vllm-workspace without .git (artifact tarball unpack). +# tests/standalone_tests/python_only_compile.sh uses merge-base(HEAD, origin/main) for +# wheels.vllm.ai; compute on the agent (full git checkout) and pass into the container. +vllm_standalone_merge_base="" +checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}" +if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then + checkout="." +fi +# Pass safe.directory per-command (-c) because buildkite runs will always fail +# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids. +if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then + vllm_standalone_merge_base="$( + git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true + )" +fi +if [[ -z "${vllm_standalone_merge_base}" ]]; then + vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}" +fi +echo "INFO: passing VLLM_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}" + MYPYTHONPATH="/vllm-workspace" container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}" @@ -501,14 +547,37 @@ if is_multi_node "$commands"; then else echo "--- Single-node job" echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" + docker_run_terminal_args=(-i) + if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then + docker_run_terminal_args+=(-t) + echo "Docker interactive stdin: enabled; TTY allocation: enabled" + else + echo "Docker interactive stdin: enabled; TTY allocation: disabled" + fi + + ulimit_core_hard=$(ulimit -H -c) + if [[ "$ulimit_core_hard" == "unlimited" ]]; then + # docker run can't pass "unlimited" to --ulimit + ulimit_core_hard="-1" + fi + # Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled + coredump_flags="--ulimit core=0:$ulimit_core_hard" + if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then + # Works around https://github.com/rocm/rocm-systems/issues/6206 + coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"' + else + echo "ROCm debug agent not enabled, coredumps are disabled in the test container." + fi docker run \ + "${docker_run_terminal_args[@]}" \ --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ $RDMA_FLAGS \ --network=host \ --shm-size=16gb \ --group-add "$render_gid" \ --rm \ + $coredump_flags \ -e HF_TOKEN \ -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \ -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \ @@ -516,6 +585,11 @@ else -e AWS_SECRET_ACCESS_KEY \ -e BUILDKITE_PARALLEL_JOB \ -e BUILDKITE_PARALLEL_JOB_COUNT \ + -e TERM \ + -e FORCE_COLOR \ + -e CLICOLOR_FORCE \ + -e PY_COLORS \ + -e PYTEST_ADDOPTS \ -v "${HF_CACHE}:${HF_MOUNT}" \ -e "HF_HOME=${HF_MOUNT}" \ -e "PYTHONPATH=${MYPYTHONPATH}" \ @@ -525,6 +599,7 @@ else -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \ -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \ -e "PYTORCH_ROCM_ARCH=" \ + -e "VLLM_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}" \ --name "${container_name}" \ "${image_name}" \ /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}" diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 35513727f16..2d11dd477ea 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -8,7 +8,7 @@ set -ex CORE_RANGE=${CORE_RANGE:-0-31} OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-31} -export CMAKE_BUILD_PARALLEL_LEVEL=16 +export CMAKE_BUILD_PARALLEL_LEVEL=32 # Setup cleanup remove_docker_container() { @@ -37,8 +37,10 @@ function cpu_tests() { 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 - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" + pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits @@ -62,7 +64,6 @@ function cpu_tests() { set -e pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs" - # basic online serving docker exec cpu-test bash -c ' set -e diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 27ec0068668..032d8e78333 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -7,10 +7,49 @@ set -euox pipefail # allow to bind to different cores CORE_RANGE=${CORE_RANGE:-48-95} NUMA_NODE=${NUMA_NODE:-1} -IMAGE_NAME="cpu-test-$NUMA_NODE" +AGENT_SLOT=${AGENT_SLOT:-} +IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}" TIMEOUT_VAL=$1 TEST_COMMAND=$2 +# Disk hygiene knobs. Reclaim space only once the Docker root filesystem crosses +# DISK_USAGE_THRESHOLD percent, and cap the shared BuildKit cache at +# BUILDKIT_CACHE_MAX so subsequent builds keep reusing the hottest layers. +DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-70} +BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-80GB} + +# Reclaim disk only when the host is under pressure. We trim (not purge) the +# shared BuildKit cache so cross-job/cross-agent reuse stays intact, and only +# touch dangling images; other agents' uniquely tagged images are left alone. +prune_if_disk_pressure() { + local docker_root disk_usage + docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true) + if [ -z "$docker_root" ]; then + return 0 + fi + disk_usage=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%') + if [ "${disk_usage:-0}" -gt "$DISK_USAGE_THRESHOLD" ]; then + echo "--- :broom: Disk usage ${disk_usage}% exceeds ${DISK_USAGE_THRESHOLD}%, reclaiming space" + docker image prune -f || true + docker builder prune -f --keep-storage="$BUILDKIT_CACHE_MAX" || true + else + echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune" + fi +} + +# Always drop this agent's image once the job ends (the default builder never +# uses it as a cache source, so removing it costs no rebuild speed), then +# reclaim space if needed. Guard every docker call with `|| true` so the trap +# never overrides the test's exit code. +cleanup() { + docker image rm -f "$IMAGE_NAME" || true + prune_if_disk_pressure +} +trap cleanup EXIT + +# Free space up front so a nearly-full host doesn't fail the build. +prune_if_disk_pressure + # building the docker image echo "--- :docker: Building Docker image" docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu . diff --git a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh new file mode 100644 index 00000000000..d59ab35db3b --- /dev/null +++ b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +set -euo pipefail + +test_suite="${1:-}" + +if [[ -z "${test_suite}" ]]; then + echo "Usage: $0 " >&2 + exit 1 +fi + +case "${test_suite}" in + example) + pip install tblib==3.1.0 + + python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager + python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE + 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 nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 + 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 TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ --block-size 64 --enforce-eager + 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 + python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 + ;; + v1) + cd tests + + pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py + pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py + pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" + 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/test_serial_utils.py + pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.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 --ignore=v1/kv_connector/unit/test_offloading_connector.py + ;; + server) + pip install av + cd tests + + pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py + pytest -v -s benchmarks/test_serve_cli.py + ;; + *) + echo "Unknown Intel test suite: ${test_suite}" >&2 + exit 1 + ;; +esac diff --git a/.buildkite/scripts/hardware_ci/run-intel-test.sh b/.buildkite/scripts/hardware_ci/run-intel-test.sh index 0eb93f5a6b3..83cde9ad16c 100755 --- a/.buildkite/scripts/hardware_ci/run-intel-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-test.sh @@ -243,8 +243,10 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head # ---- Command source selection ---- commands="" +commands_source="" if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then commands="${VLLM_TEST_COMMANDS}" + commands_source="env" echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)" elif [[ $# -gt 0 ]]; then all_yaml=true @@ -303,8 +305,12 @@ if [[ -z "$commands" ]]; then fi echo "Raw commands: $commands" -commands=$(re_quote_pytest_markers "$commands") -echo "After re-quoting: $commands" +if [[ "$commands_source" != "env" ]]; then + commands=$(re_quote_pytest_markers "$commands") + echo "After re-quoting: $commands" +else + echo "Skipping re-quoting for VLLM_TEST_COMMANDS input" +fi commands=$(apply_intel_test_overrides "$commands") echo "Final commands: $commands" @@ -354,7 +360,7 @@ export HF_TOKEN ZE_AFFINITY_MASK --ipc=host \ --privileged \ -v /dev/dri/by-path:/dev/dri/by-path \ - -v "${HOME}/.cache/huggingface:/root/.cache/huggingface" \ + -v "/data/huggingface:/root/.cache/huggingface" \ --entrypoint='' \ -e HF_TOKEN \ -e ZE_AFFINITY_MASK \ @@ -363,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK -e CMDS \ --name "${container_name}" \ "${IMAGE}" \ - bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \ + bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \ >/dev/null } 9>/tmp/docker-pull.lock diff --git a/.buildkite/scripts/hardware_ci/run-npu-test.sh b/.buildkite/scripts/hardware_ci/run-npu-test.sh index 9d33a8c0b22..b925b74c4cb 100644 --- a/.buildkite/scripts/hardware_ci/run-npu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-npu-test.sh @@ -85,7 +85,7 @@ RUN pip config set global.index-url http://cache-service-vllm.nginx-pypi-cache.s # Install for pytest to make the docker build cache layer always valid RUN --mount=type=cache,target=/root/.cache/pip \ - pip install pytest>=6.0 modelscope + pip install pytest>=6.0 'modelscope<1.38' WORKDIR /workspace/vllm diff --git a/.buildkite/scripts/install-kv-connectors.sh b/.buildkite/scripts/install-kv-connectors.sh index 34c502e6b9a..b1e024709e1 100755 --- a/.buildkite/scripts/install-kv-connectors.sh +++ b/.buildkite/scripts/install-kv-connectors.sh @@ -4,6 +4,11 @@ set -euo pipefail +if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then + uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + exit 0 +fi + REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}" uv pip install --system -r "${REQUIREMENTS_FILE}" diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index ec319aa7600..91b5c3ace1b 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -130,6 +130,22 @@ docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm docker push vllm/vllm-openai-rocm:latest-base docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base +# ---- XPU ---- + +docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu + +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 +docker push vllm/vllm-openai-xpu:latest-x86_64 +docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + +docker manifest rm vllm/vllm-openai-xpu:latest || true +docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true +docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend +docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend +docker manifest push vllm/vllm-openai-xpu:latest +docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + # ---- CPU ---- # CPU images are behind separate block steps and may not have been built. # All-or-nothing: inspect both arches first, then either publish everything diff --git a/.buildkite/scripts/rocm/build-ci-base.sh b/.buildkite/scripts/rocm/build-ci-base.sh new file mode 100755 index 00000000000..23d17e17b4d --- /dev/null +++ b/.buildkite/scripts/rocm/build-ci-base.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base. + +set -euo pipefail + +metadata_get() { + local key="$1" + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data get "${key}" 2>/dev/null || true + fi +} + +main() { + local base_refreshed="" + + base_refreshed="$(metadata_get rocm-base-refresh)" + if [[ "${base_refreshed}" == "1" ]]; then + export BASE_IMAGE + export CI_BASE_PUSH_STABLE_TAG + + BASE_IMAGE="$(metadata_get rocm-base-image)" + CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)" + CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}" + + echo "Using refreshed ROCm base image for ci_base: ${BASE_IMAGE}" + echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}" + fi + + bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps +} + +main "$@" diff --git a/.buildkite/scripts/rocm/build-test-image.sh b/.buildkite/scripts/rocm/build-test-image.sh new file mode 100755 index 00000000000..9803e20d02e --- /dev/null +++ b/.buildkite/scripts/rocm/build-test-image.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# Build the ROCm CI test image or wheel artifact. +# +# When Dockerfile.rocm_base changes, always build the full image so downstream +# ROCm tests can validate the freshly rebuilt base -> ci_base -> ci image chain. + +set -euo pipefail + +metadata_get() { + local key="$1" + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data get "${key}" 2>/dev/null || true + fi +} + +use_refreshed_base_if_present() { + local base_refreshed="" + + base_refreshed="$(metadata_get rocm-base-refresh)" + if [[ "${base_refreshed}" != "1" ]]; then + return 1 + fi + + export BASE_IMAGE + export CI_BASE_IMAGE + export IMAGE_TAG_LATEST + + BASE_IMAGE="$(metadata_get rocm-base-image)" + CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)" + IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)" + + echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}" + echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}" + if [[ -n "${IMAGE_TAG_LATEST}" ]]; then + echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}" + fi + + return 0 +} + +main() { + local base_refreshed=0 + + if use_refreshed_base_if_present; then + base_refreshed=1 + fi + + if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then + echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only" + IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts + return + fi + + bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel +} + +main "$@" diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh new file mode 100755 index 00000000000..06e3e80c967 --- /dev/null +++ b/.buildkite/scripts/rocm/refresh-base-image.sh @@ -0,0 +1,513 @@ +#!/usr/bin/env bash +# Build and publish a fresh ROCm base image when Dockerfile.rocm_base changes. +# +# Normal AMD CI builds should not pay for this path. The script no-ops unless +# docker/Dockerfile.rocm_base changed relative to the branch base, the previous +# main commit, or ROCM_BASE_REFRESH_FORCE=1 is set. + +set -euo pipefail + +DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}" +BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}" +CI_IMAGE_REPO="${ROCM_CI_IMAGE_REPO:-rocm/vllm-ci}" +BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}" +DEFAULT_ROCM_BASE_METADATA_VERSION="1" +DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}" +DEFAULT_ROCM_BASE_CONTENT_ARGS="BASE_IMAGE TRITON_BRANCH TRITON_REPO PYTORCH_BRANCH PYTORCH_REPO PYTORCH_VISION_BRANCH PYTORCH_VISION_REPO PYTORCH_AUDIO_BRANCH PYTORCH_AUDIO_REPO FA_BRANCH FA_REPO AITER_BRANCH AITER_REPO MORI_BRANCH MORI_REPO PYTORCH_ROCM_ARCH PYTHON_VERSION USE_SCCACHE" + +metadata_set() { + local key="$1" + local value="$2" + + [[ -n "${value}" ]] || return 0 + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data set "${key}" "${value}" || true + fi +} + +compute_content_hash() { + local path="" + local file="" + + for path in "$@"; do + if [[ -d "${path}" ]]; then + while IFS= read -r -d '' file; do + printf 'file:%s\n' "${file}" + sha256sum "${file}" + done < <(find "${path}" -type f -print0 | sort -z) + elif [[ -f "${path}" ]]; then + printf 'file:%s\n' "${path}" + sha256sum "${path}" + else + printf 'missing:%s\n' "${path}" + fi + done | sha256sum | cut -d' ' -f1 +} + +clean_docker_tag() { + local input="$1" + echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128 +} + +tag_component() { + local input="$1" + local max_chars="${2:-24}" + + clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}" +} + +extract_arg_default() { + local arg_name="$1" + + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + "${DOCKERFILE}" | head -1 +} + +resolve_image_digest() { + local image_ref="$1" + + docker buildx imagetools inspect "${image_ref}" 2>/dev/null \ + | sed -n -E 's/^Digest:[[:space:]]+//p' \ + | head -1 || true +} + +resolve_rocm_base_arg_value() { + local arg_name="$1" + local use_sccache="$2" + + case "${arg_name}" in + USE_SCCACHE) + printf '%s\n' "${use_sccache}" + ;; + *) + extract_arg_default "${arg_name}" + ;; + esac +} + +hash_rocm_base_arg_values() { + local use_sccache="$1" + local base_image_digest="$2" + local arg_name="" + local arg_value="" + shift 2 || true + + for arg_name in "$@"; do + [[ -n "${arg_name}" ]] || continue + arg_value=$(resolve_rocm_base_arg_value "${arg_name}" "${use_sccache}") + printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-}" + if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then + printf 'arg:%s.digest=%s\n' "${arg_name}" "${base_image_digest:-unknown}" + fi + done +} + +rocm_version_from_base_image() { + local base_image="$1" + local version="" + + version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")" + tag_component "${version:-${base_image}}" 16 +} + +git_diff_changed_base() { + local range="$1" + [[ -n "$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null)" ]] +} + +short_git_ref() { + local ref="$1" + + git rev-parse --short "${ref}" 2>/dev/null || printf '%s\n' "${ref}" +} + +extract_arg_default_from_ref() { + local ref="$1" + local arg_name="$2" + local content="" + + content="$(git show "${ref}:${DOCKERFILE}" 2>/dev/null || true)" + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + <<<"${content}" | head -1 +} + +log_arg_default_changes() { + local old_ref="$1" + local new_ref="$2" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local arg_name="" + local old_value="" + local new_value="" + local changed=0 + + echo "Changed ROCm base ARG defaults:" + for arg_name in ${content_args}; do + old_value="$(extract_arg_default_from_ref "${old_ref}" "${arg_name}")" + new_value="$(extract_arg_default_from_ref "${new_ref}" "${arg_name}")" + if [[ "${old_value}" != "${new_value}" ]]; then + echo " - ${arg_name}: ${old_value:-} -> ${new_value:-}" + changed=1 + fi + done + + if [[ "${changed}" == "0" ]]; then + echo " - none detected; Dockerfile instructions changed outside tracked ARG defaults" + fi +} + +log_arg_line_diff() { + local range="$1" + local arg_diff="" + + arg_diff="$( + git diff --unified=0 "${range}" -- "${DOCKERFILE}" 2>/dev/null \ + | awk '/^[+-][[:space:]]*ARG[[:space:]]/ && $0 !~ /^(---|\+\+\+)/ { print " " $0 }' \ + || true + )" + + if [[ -n "${arg_diff}" ]]; then + echo "Changed Dockerfile ARG lines:" + printf '%s\n' "${arg_diff}" + fi +} + +log_rocm_base_change_check() { + local context="$1" + local range="$2" + local old_ref="$3" + local old_short="" + local head_short="" + + old_short="$(short_git_ref "${old_ref}")" + head_short="$(short_git_ref HEAD)" + + echo "--- :mag: ROCm base refresh check" + echo "Context: ${context}" + echo "Dockerfile: ${DOCKERFILE}" + echo "Base revision: ${old_short}" + echo "Head revision: ${head_short}" + echo "Git diff range: ${range}" +} + +log_rocm_base_rebuild_reason() { + local context="$1" + local range="$2" + local old_ref="$3" + local changed_files="" + + log_rocm_base_change_check "${context}" "${range}" "${old_ref}" + + changed_files="$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null || true)" + echo "Changed files:" + if [[ -n "${changed_files}" ]]; then + sed 's/^/ - /' <<<"${changed_files}" + else + echo " - ${DOCKERFILE}" + fi + log_arg_default_changes "${old_ref}" HEAD + log_arg_line_diff "${range}" + echo "Decision: rebuilding ROCm base image because ${DOCKERFILE} changed." +} + +rocm_base_changed_in_range() { + local context="$1" + local range="$2" + local old_ref="$3" + + if git_diff_changed_base "${range}"; then + log_rocm_base_rebuild_reason "${context}" "${range}" "${old_ref}" + return 0 + fi + + log_rocm_base_change_check "${context}" "${range}" "${old_ref}" + echo "Decision: ROCm base refresh not required; ${DOCKERFILE} is unchanged." + return 1 +} + +rocm_base_changed() { + local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" + local base_ref="refs/remotes/origin/${base_branch}" + local merge_base="" + + if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then + echo "ROCM_BASE_REFRESH_SKIP=1 set; skipping ROCm base refresh" + return 1 + fi + + if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" == "1" ]]; then + echo "ROCM_BASE_REFRESH_FORCE=1 set; refreshing ROCm base image" + return 0 + fi + + if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "Not in a git checkout; skipping ROCm base refresh unless forced" + return 1 + fi + + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + git fetch --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true) + if [[ -z "${merge_base}" ]]; then + echo "Unable to determine merge base with PR base ${base_ref}; skipping ROCm base refresh unless forced" + return 1 + fi + if rocm_base_changed_in_range \ + "pull request build against ${base_ref}" \ + "${merge_base}...HEAD" \ + "${merge_base}"; then + return 0 + fi + elif [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] \ + && git rev-parse --verify HEAD~1 >/dev/null 2>&1; then + if rocm_base_changed_in_range \ + "stable branch build; comparing against previous ${ROCM_BASE_STABLE_BRANCH:-main} commit" \ + "HEAD~1..HEAD" \ + "HEAD~1"; then + return 0 + fi + else + git fetch --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true) + if [[ -z "${merge_base}" ]]; then + echo "Unable to determine merge base with branch base ${base_ref}; skipping ROCm base refresh unless forced" + return 1 + fi + if rocm_base_changed_in_range \ + "branch build against ${base_ref}" \ + "${merge_base}...HEAD" \ + "${merge_base}"; then + return 0 + fi + fi + + return 1 +} + +should_push_stable_tag() { + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + return 1 + fi + + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then + return 0 + fi + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then + return 1 + fi + + [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" \ + && "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] +} + +setup_builder() { + echo "--- :buildkite: Setting up buildx builder for ROCm base" + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then + docker buildx use "${BUILDER_NAME}" + else + docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use + fi + docker buildx inspect --bootstrap +} + +compute_base_content_hash() { + local use_sccache="$1" + local base_image_digest="$2" + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local -a content_paths=() + local -a content_arg_names=() + + read -r -a content_paths <<< "${content_files}" + read -r -a content_arg_names <<< "${content_args}" + + { + printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")" + printf 'dockerfile:%s\n' "${DOCKERFILE}" + printf 'resolved-build-args:\n' + hash_rocm_base_arg_values \ + "${use_sccache}" "${base_image_digest}" "${content_arg_names[@]}" + } | sha256sum | cut -d' ' -f1 +} + +build_base_image() { + local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}" + local base_hash="" + local build_date="" + local build_suffix="" + local base_image_arg="" + local base_image_digest="" + local rocm_version="" + local triton_arg="" + local pytorch_arg="" + local pytorch_vision_arg="" + local pytorch_audio_arg="" + local fa_arg="" + local aiter_arg="" + local mori_arg="" + local python_version_arg="" + local pytorch_rocm_arch_arg="" + local pytorch_branch="" + local aiter_branch="" + local dependency_summary="" + local descriptor="" + local ci_descriptor="" + local descriptive_tag="" + local stable_tag="${BASE_REPO}:base" + local ci_descriptive_tag="" + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local content_files_hash="" + local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}" + local -a tags=() + local -a no_cache_args=() + local -a sccache_args=() + local -a content_paths=() + + if [[ ! -f "${DOCKERFILE}" ]]; then + echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2 + exit 1 + fi + + build_date="${ROCM_BASE_TAG_DATE:-$(date -u +%Y%m%d)}" + if [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then + build_suffix="_bk_${BUILDKITE_BUILD_NUMBER}" + fi + base_image_arg="$(extract_arg_default BASE_IMAGE)" + base_image_digest="$(resolve_image_digest "${base_image_arg}")" + read -r -a content_paths <<< "${content_files}" + content_files_hash="$(compute_content_hash "${content_paths[@]}")" + base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}") + rocm_version="$(rocm_version_from_base_image "${base_image_arg}")" + triton_arg="$(extract_arg_default TRITON_BRANCH)" + pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)" + pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)" + pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)" + fa_arg="$(extract_arg_default FA_BRANCH)" + aiter_arg="$(extract_arg_default AITER_BRANCH)" + mori_arg="$(extract_arg_default MORI_BRANCH)" + python_version_arg="$(extract_arg_default PYTHON_VERSION)" + pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)" + pytorch_branch="$(tag_component "${pytorch_arg}" 16)" + aiter_branch="$(tag_component "${aiter_arg}" 24)" + dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}" + descriptor="$(clean_docker_tag "base_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")" + ci_descriptor="$(clean_docker_tag "ci_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")" + + descriptive_tag="${BASE_REPO}:${descriptor}" + ci_descriptive_tag="${CI_IMAGE_REPO}:${ci_descriptor}" + + tags=(-t "${descriptive_tag}") + if should_push_stable_tag; then + tags+=(-t "${stable_tag}") + metadata_set "rocm-base-push-stable-tag" "1" + else + metadata_set "rocm-base-push-stable-tag" "0" + fi + + if [[ "${ROCM_BASE_NO_CACHE:-1}" == "1" ]]; then + no_cache_args=(--no-cache) + fi + + for env_name in \ + SCCACHE_DOWNLOAD_URL \ + SCCACHE_ENDPOINT \ + SCCACHE_BUCKET_NAME \ + SCCACHE_REGION_NAME \ + SCCACHE_S3_NO_CREDENTIALS; do + if [[ -n "${!env_name:-}" ]]; then + sccache_args+=(--build-arg "${env_name}=${!env_name}") + fi + done + + echo "--- :docker: Building ROCm base image" + echo "Dockerfile: ${DOCKERFILE}" + echo "Descriptive tag: ${descriptive_tag}" + echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))" + echo "Content hash: ${base_hash}" + echo "Dependency summary: ${dependency_summary}" + echo "USE_SCCACHE: ${use_sccache}" + + docker buildx build \ + "${no_cache_args[@]}" \ + --pull \ + --progress "${BUILDKIT_PROGRESS:-plain}" \ + --file "${DOCKERFILE}" \ + --build-arg "USE_SCCACHE=${use_sccache}" \ + "${sccache_args[@]}" \ + --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \ + --label "org.opencontainers.image.vendor=vLLM" \ + --label "org.opencontainers.image.title=vLLM ROCm base" \ + --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT:-}" \ + --label "vllm.rocm_base.metadata_version=${metadata_version}" \ + --label "vllm.rocm_base.content_hash=${base_hash}" \ + --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \ + --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \ + --label "vllm.rocm_base.image.descriptive=${descriptive_tag}" \ + --label "vllm.rocm_base.image.stable=${stable_tag}" \ + --label "vllm.rocm_base.git_commit=${BUILDKITE_COMMIT:-}" \ + --label "vllm.rocm_base.stable_branch=${ROCM_BASE_STABLE_BRANCH:-main}" \ + --label "vllm.rocm_base.descriptor=${descriptor}" \ + --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \ + --label "vllm.rocm_base.base_image=${base_image_arg}" \ + --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \ + --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \ + --label "vllm.rocm_base.dependency.python=${python_version_arg}" \ + --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \ + --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \ + --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \ + --label "vllm.rocm_base.dependency.triton=${triton_arg}" \ + --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \ + --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \ + --label "vllm.rocm_base.dependency.mori=${mori_arg}" \ + --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \ + "${tags[@]}" \ + --push \ + . + + docker buildx imagetools inspect "${descriptive_tag}" >/dev/null + + metadata_set "rocm-base-refresh" "1" + metadata_set "rocm-base-image" "${descriptive_tag}" + metadata_set "rocm-base-image-descriptive" "${descriptive_tag}" + metadata_set "rocm-base-image-stable" "${stable_tag}" + metadata_set "rocm-base-image-ci-descriptive" "${ci_descriptive_tag}" + metadata_set "rocm-base-metadata-version" "${metadata_version}" + metadata_set "rocm-base-content-hash" "${base_hash}" + metadata_set "rocm-base-content-files-hash" "${content_files_hash}" + metadata_set "rocm-base-content-files" "${content_files}" + metadata_set "rocm-base-content-args" "${content_args}" + metadata_set "rocm-base-base-image-digest" "${base_image_digest}" + metadata_set "rocm-base-dockerfile" "${DOCKERFILE}" + metadata_set "rocm-base-descriptor" "${descriptor}" + metadata_set "rocm-base-dependency-summary" "${dependency_summary}" + metadata_set "rocm-base-dependency-rocm" "${rocm_version}" + metadata_set "rocm-base-dependency-python" "${python_version_arg}" + metadata_set "rocm-base-dependency-pytorch" "${pytorch_arg}" + metadata_set "rocm-base-dependency-torchvision" "${pytorch_vision_arg}" + metadata_set "rocm-base-dependency-torchaudio" "${pytorch_audio_arg}" + metadata_set "rocm-base-dependency-triton" "${triton_arg}" + metadata_set "rocm-base-dependency-flash-attention" "${fa_arg}" + metadata_set "rocm-base-dependency-aiter" "${aiter_arg}" + metadata_set "rocm-base-dependency-mori" "${mori_arg}" + metadata_set "rocm-base-pytorch-rocm-arch" "${pytorch_rocm_arch_arg}" + metadata_set "rocm-ci-image-descriptive" "${ci_descriptive_tag}" + + echo "--- :white_check_mark: ROCm base image published" + echo "Use BASE_IMAGE=${descriptive_tag} for downstream ROCm CI builds" +} + +main() { + metadata_set "rocm-base-refresh" "0" + + if ! rocm_base_changed; then + echo "ROCm base Dockerfile did not change; skipping base image refresh" + return 0 + fi + + setup_builder + build_base_image +} + +main "$@" diff --git a/.buildkite/scripts/rocm/smoke-test-image.sh b/.buildkite/scripts/rocm/smoke-test-image.sh new file mode 100755 index 00000000000..ed511c9b77a --- /dev/null +++ b/.buildkite/scripts/rocm/smoke-test-image.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Fast structural smoke test for the full ROCm CI image. + +set -euo pipefail + +image_ref="${VLLM_CI_SMOKE_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}" + +docker run --rm --network=none --entrypoint /bin/bash "${image_ref}" -ec ' + if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi + if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi + if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi + if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi + + command -v python3 + command -v uv + command -v pytest + + if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then + echo No ROCm CLI found in image >&2 + exit 1 + fi + + python3 - </dev/null 2>&1; then - return - fi - - local version="${PROTOC_VERSION:-31.1}" local arch case "$(uname -m)" in x86_64) @@ -40,16 +49,17 @@ install_protoc() { ;; esac - local url="https://github.com/protocolbuffers/protobuf/releases/download/v${version}/protoc-${version}-linux-${arch}.zip" + local url="https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-${arch}.zip" local tmp_dir tmp_dir="$(mktemp -d)" - log_section "Installing protoc ${version}" + log_section "Installing protoc ${PROTOC_VERSION}" curl -L --proto '=https' --tlsv1.2 -sSf "$url" -o "$tmp_dir/protoc.zip" mkdir -p "$CARGO_HOME/bin" unzip -q "$tmp_dir/protoc.zip" bin/protoc 'include/*' -d "$CARGO_HOME" chmod +x "$CARGO_HOME/bin/protoc" rm -rf "$tmp_dir" + protoc --version } rust_toolchain() { @@ -70,48 +80,72 @@ install_rust_toolchain() { } install_cargo_binstall() { - if command -v cargo-binstall >/dev/null 2>&1; then - return - fi - - log_section "Installing cargo-binstall" + log_section "Installing cargo-binstall ${CARGO_BINSTALL_VERSION}" curl -L --proto '=https' --tlsv1.2 -sSf \ - https://raw.githubusercontent.com/cargo-bins/cargo-binstall/main/install-from-binstall-release.sh \ - | bash + "https://raw.githubusercontent.com/cargo-bins/cargo-binstall/v${CARGO_BINSTALL_VERSION}/install-from-binstall-release.sh" \ + | env BINSTALL_VERSION="$CARGO_BINSTALL_VERSION" bash + cargo-binstall -V } install_cargo_sort() { - if command -v cargo-sort >/dev/null 2>&1; then - return - fi + log_section "Installing cargo-sort ${CARGO_SORT_VERSION_REQ}" + cargo binstall --no-confirm --force "cargo-sort@${CARGO_SORT_VERSION_REQ}" +} - log_section "Installing cargo-sort" - install_cargo_binstall - cargo binstall --no-confirm cargo-sort +install_cargo_deny() { + log_section "Installing cargo-deny ${CARGO_DENY_VERSION_REQ}" + cargo binstall --no-confirm --force "cargo-deny@${CARGO_DENY_VERSION_REQ}" } install_cargo_nextest() { - if command -v cargo-nextest >/dev/null 2>&1; then - return - fi - - log_section "Installing cargo-nextest" - install_cargo_binstall - cargo binstall --no-confirm --secure cargo-nextest + log_section "Installing cargo-nextest ${CARGO_NEXTEST_VERSION_REQ}" + cargo binstall \ + --no-confirm \ + --force \ + --secure \ + "cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}" } install_uv() { - if command -v uv >/dev/null 2>&1; then - return - fi - - log_section "Installing uv" - curl -LsSf --proto '=https' --tlsv1.2 https://astral.sh/uv/install.sh \ + log_section "Installing uv ${UV_VERSION}" + curl -L --proto '=https' --tlsv1.2 -sSf \ + "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-installer.sh" \ | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh + uv --version +} + +setup_pyo3_python() { + log_section "Installing Python ${PYO3_PYTHON_VERSION} for PyO3 tests" + uv python install "$PYO3_PYTHON_VERSION" + PYO3_PYTHON="$(uv python find \ + --managed-python \ + --no-project \ + --resolve-links \ + "$PYO3_PYTHON_VERSION")" + export PYO3_PYTHON + + local python_libdir + python_libdir="$("$PYO3_PYTHON" - <<'PY' +import pathlib +import sysconfig + +libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR")) +ldlibrary = sysconfig.get_config_var("LDLIBRARY") +assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1 +assert ldlibrary +assert (libdir / ldlibrary).exists(), libdir / ldlibrary +print(libdir) +PY +)" + + export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}" + export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}" } run_style_clippy() { + install_cargo_binstall install_cargo_sort + install_cargo_deny log_section "Checking Rust formatting" cargo fmt --manifest-path rust/Cargo.toml --all -- --check @@ -119,6 +153,13 @@ run_style_clippy() { log_section "Checking Cargo.toml ordering" cargo sort --workspace --check rust + log_section "Checking Rust dependency bans" + cargo deny \ + --manifest-path rust/Cargo.toml \ + --config rust/deny.toml \ + check \ + bans + log_section "Running clippy" cargo clippy \ --manifest-path rust/Cargo.toml \ @@ -132,6 +173,8 @@ run_style_clippy() { run_tests() { install_uv + setup_pyo3_python + install_cargo_binstall install_cargo_nextest log_section "Running cargo nextest" diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh index 0eadfa1f80b..e1808835fdf 100755 --- a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh @@ -33,6 +33,14 @@ if [[ -n "${ATTENTION_BACKEND:-}" ]]; then EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}") fi +# ROCm: run eager to avoid intermittent HIP-graph decode corruption. +# See https://github.com/ROCm/clr/issues/279 +# TODO(aarushjain29): Revert after TheRock 7.14 +if command -v rocm-smi &> /dev/null || command -v amd-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then + echo "ROCm platform detected: adding --enforce-eager to avoid HIP-graph decode corruption" + EXTRA_ARGS+=(--enforce-eager) +fi + cleanup() { if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then kill "${SERVER_PID}" 2>/dev/null || true 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 index 06743f16b68..82d4e27c219 100755 --- 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 @@ -18,6 +18,10 @@ wait_for_server() { MODEL="Qwen/Qwen3-30B-A3B-FP8" BACK="allgather_reducescatter" +if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then + # Disable MOE padding for ROCm since it is causing eplb to fail. + export VLLM_ROCM_MOE_PADDING=0 +fi cleanup() { if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then diff --git a/.buildkite/scripts/trigger-ci-build.sh b/.buildkite/scripts/trigger-ci-build.sh index 65a761588c9..a56683e9603 100755 --- a/.buildkite/scripts/trigger-ci-build.sh +++ b/.buildkite/scripts/trigger-ci-build.sh @@ -23,6 +23,7 @@ NC='\033[0m' # No Color # Default configuration PIPELINE="ci" DRY_RUN=true +TORCH_NIGHTLY=false usage() { cat <= 0.0.2" + - pytest -v -s plugins_tests/gguf + +#------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# + +- label: Rust Frontend OpenAI Coverage # 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: + - rust/ + - vllm/benchmarks/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/sample/ + - tests/utils.py + - tests/benchmarks/test_serve_cli.py + - tests/entrypoints/openai/chat_completion/test_chat_completion.py + - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py + - tests/entrypoints/openai/completion/test_shutdown.py + - tests/entrypoints/openai/test_return_token_ids.py + - tests/entrypoints/openai/test_uds.py + - tests/v1/sample/test_logprobs_e2e.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" + - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" + - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" + - pytest -v -s entrypoints/openai/test_uds.py + - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" + +- label: Rust Frontend Serve Admin Coverage # 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: + - rust/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/engine/ + - tests/utils.py + - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py + - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py + - tests/entrypoints/serve/instrumentator/test_basic.py + - tests/entrypoints/serve/instrumentator/test_metrics.py + - tests/entrypoints/serve/tokenize/test_tokenization.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" + - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" + +- label: Rust Frontend Core Correctness # 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: + - rust/ + - vllm/entrypoints/openai/ + - tests/utils.py + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: Rust Frontend Tool Use # 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: + - rust/ + - vllm/entrypoints/openai/ + - vllm/tool_parsers/ + - tests/utils.py + - tests/tool_use/ + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" + +- label: Rust Frontend Distributed # 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: + - rust/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/utils.py + - tests/v1/distributed/test_external_lb_dp.py + - tests/v1/distributed/test_hybrid_lb_dp.py + - tests/v1/distributed/test_internal_lb_dp.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" + #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# - label: Quantization # TBD @@ -2040,7 +2176,7 @@ steps: - vllm/ - tests/test_regression commands: - - pip install modelscope + - pip install 'modelscope<1.38' - pytest -v -s test_regression.py #--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# @@ -2074,19 +2210,6 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test -- label: e2e Core (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/v1/ - - tests/v1/e2e/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py - - label: e2e Scheduling (1 GPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2135,6 +2258,7 @@ steps: 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/v1/spec_decode/ @@ -2216,7 +2340,8 @@ steps: - label: Extract Hidden States Integration # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - agent_pool: mi300_1 + agent_pool: mi300_2 + num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/config/speculative.py @@ -2282,7 +2407,7 @@ steps: - 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 + - pip install -U git+https://github.com/vllm-project/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 @@ -2327,6 +2452,7 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 + optional: true num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2345,30 +2471,11 @@ 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 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/" - source_file_dependencies: - - vllm/distributed/ - - vllm/v1/distributed/ - - vllm/model_executor/layers/fused_moe/ - - tests/v1/distributed/test_dbo.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - - 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: Metrics, Tracing (2 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 + optional: true num_gpus: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2394,11 +2501,28 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" +- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # 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/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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 + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh + - 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 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -2437,7 +2561,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: @@ -2449,9 +2573,10 @@ steps: 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 + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: @@ -2463,6 +2588,7 @@ steps: 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 @@ -2472,6 +2598,59 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # 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/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh + +- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # 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/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh + +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # 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/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + - label: V1 e2e (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2640,27 +2819,6 @@ 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: ROCm LM Eval Large Models (8 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_8 - optional: true - num_gpus: 8 - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - 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 - commands: - - 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 / language -----------------------------------------------------# - label: Language Models Test (Extended Generation) # TBD @@ -2680,7 +2838,6 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_1 - torch_nightly: true parallelism: 2 optional: true working_dir: "/vllm-workspace/tests" @@ -2710,14 +2867,12 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_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 @@ -2754,7 +2909,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py 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 + - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" #-------------------------------------------------------- mi355 · distributed --------------------------------------------------------# @@ -2786,15 +2941,19 @@ steps: - vllm/envs.py - examples/offline_inference/data_parallel.py - tests/distributed/test_context_parallel.py + - tests/distributed/test_rocm_aiter_custom_ar.py - tests/distributed/test_rocm_quick_reduce.py - tests/distributed/test_quick_all_reduce.py + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/v1/distributed/test_dbo.py - tests/utils.py commands: - pytest -v -s tests/distributed/test_context_parallel.py - - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py + - pytest -v -s tests/v1/distributed/test_dbo.py #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# @@ -2804,22 +2963,22 @@ steps: agent_pool: mi355_1 optional: true fast_check: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - tests/entrypoints/serve + - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc + - pytest -v -s entrypoints/scale_out - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2835,7 +2994,6 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2852,7 +3010,6 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2873,7 +3030,6 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2887,7 +3043,6 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2901,7 +3056,6 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 fast_check: true - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2935,7 +3089,7 @@ steps: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt - label: LM Eval Qwen3-5 Models (B200-MI355) # TBD - timeout_in_minutes: 120 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 num_gpus: 2 @@ -2946,7 +3100,6 @@ steps: - 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 @@ -3030,7 +3183,7 @@ steps: commands: - pip install tensorizer # Basic - - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN + - python3 basic/offline_inference/chat.py - 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 @@ -3075,13 +3228,13 @@ steps: - rocm-smi - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN - pytest -v -s tests/kernels/attention/test_attention_selector.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 parallelism: 2 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ @@ -3098,7 +3251,7 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - parallelism: 4 + parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ @@ -3184,7 +3337,6 @@ steps: - 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 @@ -3206,7 +3358,6 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 - torch_nightly: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -3228,7 +3379,6 @@ steps: - 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 @@ -3242,7 +3392,6 @@ steps: - 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 (Extended Pooling) # TBD @@ -3261,14 +3410,12 @@ steps: 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 @@ -3276,14 +3423,12 @@ steps: 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 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 @@ -3367,7 +3512,7 @@ steps: - pytest -v -s v1/attention - label: V1 Core + KV + Metrics # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 optional: true @@ -3389,11 +3534,11 @@ steps: - 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 + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - label: V1 Sample + Logits # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 optional: true @@ -3413,7 +3558,7 @@ steps: - pytest -v -s v1/test_outputs.py - label: V1 Spec Decode # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" @@ -3423,52 +3568,6 @@ steps: commands: - pytest -v -s -m 'not slow_test' v1/spec_decode -- 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: - - 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 - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh - -- 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/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 - - ATTENTION_BACKEND=TRITON_ATTN 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, amdgfx950nightly, amdmi355] - agent_pool: mi355_4 - num_gpus: 4 - optional: true - 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 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - #------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# - label: Weight Loading Multiple GPU # TBD @@ -3508,5 +3607,5 @@ steps: - vllm/ - tests/test_regression commands: - - pip install modelscope + - pip install 'modelscope<1.38' - pytest -v -s test_regression.py diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 8814a3a8f0c..aadea2908c6 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: V1 attention (H100-MI300) key: v1-attention-h100-mi300 - timeout_in_minutes: 30 + timeout_in_minutes: 85 device: h100 source_file_dependencies: - vllm/config/attention.py @@ -12,11 +12,12 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + parallelism: 2 mirror: amd: - device: mi300_1 - timeout_in_minutes: 70 + device: mi325_1 + timeout_in_minutes: 95 depends_on: - image-build-amd source_file_dependencies: @@ -30,7 +31,7 @@ steps: - label: V1 attention (B200) key: v1-attention-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 80 device: b200-k8s source_file_dependencies: - vllm/config/attention.py @@ -38,4 +39,5 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + parallelism: 2 diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index 0310945b086..1b92babb3ce 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Basic Correctness key: basic-correctness - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/ @@ -16,3 +16,9 @@ steps: - pytest -v -s basic_correctness/test_mem.py - pytest -v -s basic_correctness/test_basic_correctness.py - pytest -v -s basic_correctness/test_cpu_offload.py + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 70 + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 85f80478017..7c26a8d6e98 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -4,13 +4,18 @@ depends_on: steps: - label: Benchmarks CLI Test key: benchmarks-cli-test - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/ - tests/benchmarks/ commands: - pytest -v -s benchmarks/ + mirror: + amd: + device: mi300_1 + depends_on: + - image-build-amd - label: Attention Benchmarks Smoke Test (B200) key: attention-benchmarks-smoke-test-b200 @@ -18,9 +23,9 @@ steps: num_gpus: 2 optional: true working_dir: "/vllm-workspace/" - timeout_in_minutes: 10 + timeout_in_minutes: 20 source_file_dependencies: - benchmarks/attention_benchmarks/ - vllm/v1/attention/ commands: - - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1 + - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 01248738d51..eba37f657d9 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Sequence Parallel Correctness Tests (2 GPUs) key: sequence-parallel-correctness-tests-2-gpus - timeout_in_minutes: 50 + timeout_in_minutes: 80 working_dir: "/vllm-workspace/" num_devices: 2 source_file_dependencies: @@ -19,7 +19,7 @@ steps: - label: Sequence Parallel Correctness Tests (2xH100) key: sequence-parallel-correctness-tests-2xh100 - timeout_in_minutes: 50 + timeout_in_minutes: 75 working_dir: "/vllm-workspace/" device: h100 optional: true @@ -30,7 +30,7 @@ steps: - label: AsyncTP Correctness Tests (2xH100) key: asynctp-correctness-tests-2xh100 - timeout_in_minutes: 50 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: h100 optional: true @@ -41,7 +41,7 @@ steps: - label: AsyncTP Correctness Tests (B200) key: asynctp-correctness-tests-b200 - timeout_in_minutes: 50 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s optional: true @@ -52,7 +52,7 @@ steps: - label: Distributed Compile Unit Tests (2xH100) key: distributed-compile-unit-tests-2xh100 - timeout_in_minutes: 20 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -66,7 +66,7 @@ steps: - label: Fusion and Compile Unit Tests (2xB200) key: fusion-and-compile-unit-tests-2xb200 - timeout_in_minutes: 20 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s source_file_dependencies: @@ -96,7 +96,7 @@ steps: - label: Fusion E2E Quick (H100) key: fusion-e2e-quick-h100 - timeout_in_minutes: 15 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 num_devices: 1 @@ -115,7 +115,7 @@ steps: - label: Fusion E2E Config Sweep (H100) key: fusion-e2e-config-sweep-h100 - timeout_in_minutes: 30 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 num_devices: 1 @@ -149,7 +149,7 @@ steps: - label: Fusion E2E TP2 Quick (H100) key: fusion-e2e-tp2-quick-h100 - timeout_in_minutes: 20 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -167,7 +167,7 @@ steps: - label: Fusion E2E TP2 AR-RMS Config Sweep (H100) key: fusion-e2e-tp2-ar-rms-config-sweep-h100 - timeout_in_minutes: 40 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -207,7 +207,7 @@ steps: - label: Fusion E2E TP2 (B200) key: fusion-e2e-tp2-b200 - timeout_in_minutes: 20 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: b200-k8s num_devices: 2 diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index b56e635bea6..99e1949fef7 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -2,9 +2,9 @@ group: CUDA depends_on: - image-build steps: -- label: Platform Tests (CUDA) - key: platform-tests-cuda - timeout_in_minutes: 15 +- label: Platform Tests + key: platform-tests + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/envs.py @@ -19,7 +19,7 @@ steps: - label: Cudagraph key: cudagraph - timeout_in_minutes: 20 + timeout_in_minutes: 30 source_file_dependencies: - tests/v1/cudagraph - vllm/v1/cudagraph_dispatcher.py diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index c9d5237b67b..38020d91ad0 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed NixlConnector PD accuracy (4 GPUs) key: distributed-nixlconnector-pd-accuracy-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -13,9 +13,23 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + mirror: + amd: + device: mi300_4 + timeout_in_minutes: 85 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs) key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -25,6 +39,19 @@ steps: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: Push NixlConnector PP prefill PD accuracy (4 GPUs) + key: push-nixlconnector-pp-prefill-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/ + - tests/v1/kv_connector/nixl_integration/ + - tests/v1/kv_connector/nixl_push_integration/ + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh + - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus timeout_in_minutes: 30 @@ -36,10 +63,23 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + mirror: + amd: + device: mi300_4 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -48,10 +88,23 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + mirror: + amd: + device: mi300_4 + timeout_in_minutes: 85 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus - timeout_in_minutes: 25 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -60,10 +113,37 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + mirror: + amd: + device: mi300_4 + timeout_in_minutes: 80 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + +- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) + key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus + timeout_in_minutes: 25 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) key: multiconnector-nixl-offloading-pd-accuracy-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -78,7 +158,7 @@ steps: - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) key: nixlconnector-pd-spec-decode-acceptance-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: a100 working_dir: "/vllm-workspace/tests" num_devices: 2 @@ -89,10 +169,24 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh + mirror: + amd: + device: mi300_2 + timeout_in_minutes: 70 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - 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 + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 8aa41a9a26a..b519132c6f1 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed Comm Ops key: distributed-comm-ops - timeout_in_minutes: 20 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -18,7 +18,7 @@ steps: - label: Distributed DP Tests (2 GPUs) key: distributed-dp-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -37,10 +37,25 @@ steps: - 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 + mirror: + amd: + device: mi300_2 + depends_on: + - image-build-amd + 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/platforms/rocm.py - label: Distributed Compile + RPC Tests (2 GPUs) key: distributed-compile-rpc-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 65 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -63,7 +78,7 @@ steps: - label: Distributed Torchrun + Shutdown Tests (2 GPUs) key: distributed-torchrun-shutdown-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -118,7 +133,7 @@ steps: - label: Distributed DP Tests (4 GPUs) key: distributed-dp-tests-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -139,7 +154,7 @@ steps: - label: Distributed Compile + Comm (4 GPUs) key: distributed-compile-comm-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 70 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -159,9 +174,9 @@ steps: # test multi-node TP with multiproc executor (simulated on single node) - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node -- label: Distributed Tests (8 GPUs)(H100) - key: distributed-tests-8-gpus-h100 - timeout_in_minutes: 10 +- label: Distributed Tests (8xH100) + key: distributed-tests-8xh100 + timeout_in_minutes: 20 device: h100 num_devices: 8 working_dir: "/vllm-workspace/tests" @@ -180,8 +195,8 @@ steps: # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep -- label: Distributed Tests (4 GPUs)(A100) - key: distributed-tests-4-gpus-a100 +- label: Distributed Tests (4xA100) + key: distributed-tests-4xa100 device: a100 optional: true num_devices: 4 @@ -195,9 +210,9 @@ steps: - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py -- label: Distributed Tests (2 GPUs)(H100) - key: distributed-tests-2-gpus-h100 - timeout_in_minutes: 15 +- label: Distributed Tests (2xH100-2xMI300) + key: distributed-tests-2xh100-2xmi300 + timeout_in_minutes: 30 device: h100 optional: true working_dir: "/vllm-workspace/" @@ -210,15 +225,15 @@ steps: - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py - pytest -v -s tests/distributed/test_packed_tensor.py -- label: Distributed Tests (2 GPUs)(B200) - key: distributed-tests-2-gpus-b200 +- label: Distributed Tests (2xB200) + key: distributed-tests-2xb200 device: b200-k8s optional: true working_dir: "/vllm-workspace/" num_devices: 2 commands: - pytest -v -s tests/distributed/test_context_parallel.py - - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py + - pytest -v -s tests/distributed/test_nccl_symm_mem.py - pytest -v -s tests/v1/distributed/test_dbo.py - pytest -v -s tests/distributed/test_mnnvl_alltoall.py @@ -244,7 +259,7 @@ steps: - label: Pipeline + Context Parallelism (4 GPUs) key: pipeline-context-parallelism-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -259,7 +274,7 @@ steps: - label: RayExecutorV2 (4 GPUs) key: rayexecutorv2-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 9bf96221abe..9f0562ca3bc 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -3,7 +3,7 @@ depends_on: - image-build-cpu steps: - label: Docker Build Metadata - timeout_in_minutes: 10 + timeout_in_minutes: 20 device: cpu-small source_file_dependencies: - .buildkite/release-pipeline.yaml diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 88039a33960..6655ae781e8 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -2,9 +2,9 @@ group: E2E Integration depends_on: - image-build steps: -- label: DeepSeek V2-Lite Sync EPLB Accuracy - key: deepseek-v2-lite-sync-eplb-accuracy - timeout_in_minutes: 60 +- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100) + key: deepseek-v2-lite-sync-eplb-accuracy-4xh100 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -12,9 +12,9 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy - key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy - timeout_in_minutes: 60 +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100) + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -22,9 +22,9 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200) - key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200 - timeout_in_minutes: 60 +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200) + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200 + timeout_in_minutes: 20 device: b200-k8s optional: true num_devices: 2 @@ -34,7 +34,7 @@ steps: - label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -44,7 +44,7 @@ steps: - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) key: deepseek-v2-lite-prefetch-offload-accuracy-h100 - timeout_in_minutes: 60 + timeout_in_minutes: 20 device: h100 optional: true num_devices: 1 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index ddae5e774f6..1c50ab7a113 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Engine key: engine - timeout_in_minutes: 15 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/compilation/ @@ -28,14 +28,14 @@ steps: - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py mirror: amd: - device: mi300_1 - timeout_in_minutes: 60 + device: mi325_1 + timeout_in_minutes: 50 depends_on: - image-build-amd - label: Engine (1 GPU) key: engine-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 45 source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ @@ -44,14 +44,14 @@ steps: - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py mirror: amd: - device: mi300_1 - timeout_in_minutes: 40 + device: mi325_1 + timeout_in_minutes: 55 depends_on: - image-build-amd - label: e2e Scheduling (1 GPU) key: e2e-scheduling-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 35 device: h200_18gb source_file_dependencies: - vllm/v1/ @@ -60,24 +60,34 @@ steps: - pytest -v -s v1/e2e/general/test_async_scheduling.py mirror: amd: - device: mi250_1 - timeout_in_minutes: 60 + device: mi325_1 + timeout_in_minutes: 70 depends_on: - image-build-amd - label: e2e Core (1 GPU) device: h200_35gb key: e2e-core-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ commands: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/ + - tests/v1/e2e/general/ + - vllm/platforms/rocm.py - label: V1 e2e (2 GPUs) key: v1-e2e-2-gpus - timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability + timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability optional: true num_devices: 2 source_file_dependencies: @@ -102,10 +112,15 @@ steps: commands: # Only run tests that need exactly 2 GPUs - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism" + mirror: + amd: + device: mi300_2 + depends_on: + - image-build-amd - label: V1 e2e (4 GPUs) key: v1-e2e-4-gpus - timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability + timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability optional: true num_devices: 4 source_file_dependencies: @@ -133,7 +148,7 @@ steps: - label: V1 e2e (4xH100) key: v1-e2e-4xh100 - timeout_in_minutes: 60 + timeout_in_minutes: 35 device: h100 num_devices: 4 optional: true diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 4dc6d7b0e81..2db4c5ad5a2 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Entrypoints Unit Tests key: entrypoints-unit-tests - timeout_in_minutes: 10 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/entrypoints @@ -16,7 +16,7 @@ steps: - label: Entrypoints Integration (LLM) key: entrypoints-integration-llm - timeout_in_minutes: 40 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -29,30 +29,34 @@ steps: mirror: amd: device: mi325_1 + # TODO(akaratza): Test after Torch >= 2.12 bump + soft_fail: true depends_on: - image-build-amd - label: Entrypoints Integration (API Server) key: entrypoints-integration-api-server device: h200_35gb - timeout_in_minutes: 130 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - tests/entrypoints/serve + - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc + - pytest -v -s entrypoints/scale_out mirror: amd: - device: mi300_1 + device: mi325_1 depends_on: - image-build-amd - label: Entrypoints Integration (API Server OpenAI - Part 1) key: entrypoints-integration-api-server-openai-part-1 - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -63,14 +67,14 @@ steps: - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness mirror: amd: - device: mi300_1 - timeout_in_minutes: 80 + device: mi325_1 + timeout_in_minutes: 65 depends_on: - image-build-amd - label: Entrypoints Integration (API Server OpenAI - Part 2) key: entrypoints-integration-api-server-openai-part-2 - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -82,7 +86,7 @@ steps: - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py mirror: amd: - device: mi300_1 + device: mi325_1 timeout_in_minutes: 80 depends_on: - image-build-amd @@ -104,8 +108,8 @@ steps: - pytest -v -s entrypoints/anthropic mirror: amd: - device: mi300_1 - timeout_in_minutes: 60 + device: mi325_1 + timeout_in_minutes: 65 depends_on: - image-build-amd @@ -122,7 +126,7 @@ steps: - label: Entrypoints Integration (Speech to Text) device: h200_35gb key: entrypoints-integration-speech_to_text - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -134,7 +138,7 @@ steps: - label: Entrypoints Integration (Multimodal) device: h200_35gb key: entrypoints-integration-multimodal - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -156,7 +160,7 @@ steps: - label: OpenAI API Correctness key: openai-api-correctness - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - csrc/ @@ -165,7 +169,7 @@ steps: - pytest -s entrypoints/openai/correctness/ mirror: amd: - device: mi300_1 + device: mi325_1 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 0f7ab0d7157..d02ffeb7748 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: EPLB Algorithm key: eplb-algorithm - timeout_in_minutes: 15 + timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -14,10 +14,20 @@ steps: commands: - pytest -v -s distributed/test_eplb_algo.py - pytest -v -s distributed/test_eplb_utils.py + mirror: + amd: + device: mi300_1 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/eplb + - tests/distributed/test_eplb_algo.py + - tests/distributed/test_eplb_utils.py + - vllm/platforms/rocm.py - label: EPLB Execution # 17min key: eplb-execution - timeout_in_minutes: 27 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -29,7 +39,7 @@ steps: - label: Elastic EP Scaling Test key: elastic-ep-scaling-test - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h100 working_dir: "/vllm-workspace/tests" num_devices: 4 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d1a4ade2a77..894cda34119 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: vLLM IR Tests key: vllm-ir-tests - timeout_in_minutes: 10 + timeout_in_minutes: 35 device: h200_18gb working_dir: "/vllm-workspace/" source_file_dependencies: @@ -16,18 +16,19 @@ steps: - label: Kernels Core Operation Test key: kernels-core-operation-test - timeout_in_minutes: 75 + timeout_in_minutes: 120 source_file_dependencies: - csrc/ - tests/kernels/core - tests/kernels/test_concat_mla_q.py - tests/kernels/test_fused_qk_norm_rope_gate.py commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + parallelism: 3 - label: Kernels MiniMax Reduce RMS Test (2 GPUs) key: kernels-minimax-reduce-rms-test-2-gpus - timeout_in_minutes: 15 + timeout_in_minutes: 20 num_devices: 2 device: h100 source_file_dependencies: @@ -41,18 +42,20 @@ steps: - label: Deepseek V4 Kernel Test (H100) key: deepseek-v4-kernel-test-h100 - timeout_in_minutes: 15 + timeout_in_minutes: 30 device: h100 source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s kernels/test_top_k_per_row.py - label: Deepseek V4 Kernel Test (B200) key: deepseek-v4-kernel-test-b200 - timeout_in_minutes: 15 + timeout_in_minutes: 20 device: b200-k8s source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -63,7 +66,7 @@ steps: - label: Kernels Attention Test %N key: kernels-attention-test - timeout_in_minutes: 35 + timeout_in_minutes: 65 source_file_dependencies: - csrc/attention/ - vllm/v1/attention @@ -74,10 +77,37 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 90 + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/attention/ + - vllm/v1/attention + - vllm/model_executor/layers/attention + - tests/kernels/attention + - vllm/_aiter_ops.py + - vllm/envs.py + - vllm/platforms/rocm.py + +- label: Kernels Attention DiffKV Test (H100) + key: kernels-attention-diffkv-test-h100 + timeout_in_minutes: 20 + device: h100 + num_devices: 1 + source_file_dependencies: + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py + - vllm/v1/attention/backends/triton_attn_diffkv.py + - vllm/v1/attention/backends/flash_attn_diffkv.py + - tests/kernels/attention/test_triton_unified_attention_diffkv.py + commands: + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py - label: Kernels Quantization Test %N key: kernels-quantization-test - timeout_in_minutes: 90 + timeout_in_minutes: 60 source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization @@ -87,10 +117,11 @@ steps: parallelism: 2 mirror: amd: - device: mi300_1 + device: mi325_1 source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization + - vllm/config/ - tests/kernels/quantization - tests/kernels/quantization/test_rocm_skinny_gemms.py - vllm/_aiter_ops.py @@ -101,7 +132,7 @@ steps: - label: Kernels MoE Test %N key: kernels-moe-test - timeout_in_minutes: 25 + timeout_in_minutes: 50 source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ @@ -114,10 +145,26 @@ 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 parallelism: 5 + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 65 + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: Kernels Mamba Test key: kernels-mamba-test - timeout_in_minutes: 45 + timeout_in_minutes: 40 source_file_dependencies: - csrc/mamba/ - tests/kernels/mamba @@ -126,7 +173,7 @@ steps: - pytest -v -s kernels/mamba - label: Kernels KDA Test - timeout_in_minutes: 20 + timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - vllm/model_executor/layers/fla/ops/kda.py @@ -138,7 +185,7 @@ steps: - label: Kernels DeepGEMM Test (H100) key: kernels-deepgemm-test-h100 - timeout_in_minutes: 45 + timeout_in_minutes: 35 device: h100 num_devices: 1 source_file_dependencies: @@ -165,7 +212,7 @@ steps: - label: Kernels (B200) key: kernels-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 80 working_dir: "/vllm-workspace/" device: b200-k8s # optional: true @@ -218,19 +265,20 @@ steps: - label: Kernels Helion Test key: kernels-helion-test - timeout_in_minutes: 30 + timeout_in_minutes: 115 device: h100 source_file_dependencies: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==1.0.0 - - pytest -v -s kernels/helion/ + - pip install helion==1.1.0 + - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + parallelism: 2 -- label: Kernels FP8 MoE Test (1 H100) - key: kernels-fp8-moe-test-1-h100 - timeout_in_minutes: 90 +- label: Kernels FP8 MoE Test (1xH100) + key: kernels-fp8-moe-test-1xh100 + timeout_in_minutes: 40 device: h100 num_devices: 1 optional: true @@ -245,9 +293,9 @@ steps: - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py -- label: Kernels FP8 MoE Test (2 H100s) - key: kernels-fp8-moe-test-2-h100s - timeout_in_minutes: 90 +- label: Kernels FP8 MoE Test (2xH100) + key: kernels-fp8-moe-test-2xh100 + timeout_in_minutes: 45 device: h100 num_devices: 2 optional: true @@ -257,7 +305,7 @@ steps: - label: Kernels Fp4 MoE Test (B200) key: kernels-fp4-moe-test-b200 - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: b200-k8s num_devices: 1 optional: true @@ -270,7 +318,7 @@ steps: - label: Kernels FusedMoE Layer Test (2 H100s) key: kernels-fusedmoe-layer-test-2-h100s - timeout_in_minutes: 90 + timeout_in_minutes: 30 device: h100 num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 8219fa19155..1f7d2787536 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -5,7 +5,7 @@ steps: - label: LM Eval Small Models device: h200_35gb key: lm-eval-small-models - timeout_in_minutes: 75 + timeout_in_minutes: 45 source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -14,7 +14,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt mirror: amd: - device: mi300_1 + device: mi325_1 timeout_in_minutes: 55 depends_on: - image-build-amd @@ -28,7 +28,8 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py -# - label: LM Eval Large Models (4 GPUs)(A100) +# - label: LM Eval Large Models (4xA100) +# key: lm-eval-large-models-4xa100 # device: a100 # optional: true # num_devices: 4 @@ -40,8 +41,8 @@ steps: # - 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: LM Eval Large Models (4 GPUs)(H100) - key: lm-eval-large-models-4-gpus-h100 +- label: LM Eval Large Models (4xH100) + key: lm-eval-large-models-4xh100 device: h100 optional: true num_devices: 4 @@ -53,9 +54,9 @@ steps: - export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 -- label: LM Eval Small Models (B200) - key: lm-eval-small-models-b200 - timeout_in_minutes: 120 +- label: LM Eval Small Models (1xB200) + key: lm-eval-small-models-1xb200 + timeout_in_minutes: 50 device: b200-k8s optional: true source_file_dependencies: @@ -64,10 +65,23 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt -- label: LM Eval Large Models (B200, EP) - key: lm-eval-large-models-b200-ep +- label: LM Eval Small Models Distributed (2xB200) + key: lm-eval-small-models-distributed-2xb200 timeout_in_minutes: 120 device: b200-k8s + num_devices: 2 + optional: true + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt + +- label: LM Eval Large Models EP (2xB200) + key: lm-eval-large-models-ep-2xb200 + timeout_in_minutes: 60 + device: b200-k8s optional: true num_devices: 2 source_file_dependencies: @@ -76,9 +90,9 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt -- label: LM Eval Qwen3.5 Models (B200) - key: lm-eval-qwen3-5-models-b200 - timeout_in_minutes: 120 +- label: LM Eval Qwen3.5 Models (2xB200) + key: lm-eval-qwen3-5-models-2xb200 + timeout_in_minutes: 45 device: b200-k8s optional: true num_devices: 2 @@ -93,14 +107,24 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt -- label: LM Eval Large Models (H200) - key: lm-eval-large-models-h200 - timeout_in_minutes: 60 +- label: LM Eval Large Models (8xH200) + key: lm-eval-large-models-8xh200 + timeout_in_minutes: 50 device: h200 optional: true num_devices: 8 commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt + mirror: + amd: + device: mi300_8 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - label: MoE Refactor Integration Test (H100 - TEMPORARY) key: moe-refactor-integration-test-h100-temporary @@ -126,10 +150,101 @@ steps: commands: - 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 Humming f16 (A100 - TEMPORARY) + key: lm-eval-humming-f16-a100 + timeout_in_minutes: 75 + device: a100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + +- label: LM Eval Humming Act int8 (A100 - TEMPORARY) + key: lm-eval-humming-act-a100 + timeout_in_minutes: 45 + device: a100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (H100 - TEMPORARY) + key: lm-eval-humming-f16-h100 + timeout_in_minutes: 70 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + +- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) + key: lm-eval-humming-act-h100 + timeout_in_minutes: 70 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (B200 - TEMPORARY) + key: lm-eval-humming-f16-b200 + timeout_in_minutes: 50 + device: b200-k8s + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + +- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY) + key: lm-eval-humming-act-b200 + timeout_in_minutes: 50 + device: b200-k8s + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt - label: LM Eval TurboQuant KV Cache key: lm-eval-turboquant-kv-cache - timeout_in_minutes: 75 + timeout_in_minutes: 55 device: h200_18gb source_file_dependencies: - vllm/model_executor/layers/quantization/turboquant/ @@ -139,9 +254,9 @@ steps: 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) - key: gpqa-eval-gpt-oss-h100 - timeout_in_minutes: 120 +- label: GPQA Eval (GPT-OSS) (2xH100) + key: gpqa-eval-gpt-oss-2xh100 + timeout_in_minutes: 35 device: h100 optional: true num_devices: 2 @@ -153,9 +268,9 @@ steps: - 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-h100.txt -- label: GPQA Eval (GPT-OSS) (B200) - key: gpqa-eval-gpt-oss-b200 - timeout_in_minutes: 120 +- label: GPQA Eval (GPT-OSS) (2xB200) + key: gpqa-eval-gpt-oss-2xb200 + timeout_in_minutes: 30 device: b200-k8s optional: true num_devices: 2 @@ -167,9 +282,66 @@ steps: - 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-b200.txt +- label: GPQA Eval (GPT-OSS) (DGX Spark) + key: gpqa-eval-gpt-oss-spark + timeout_in_minutes: 35 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/evals/gpt_oss/ + 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-spark.txt + +- label: LM Eval KV-Offload (1xH200) + key: kv-offload-small + timeout_in_minutes: 30 + device: h200_35gb + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" + +- label: LM Eval KV-Offload (2xH100) + key: kv-offload-medium + timeout_in_minutes: 30 + device: h100 + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b" + +- label: LM Eval KV-Offload (4xH100) + key: kv-offload-large + timeout_in_minutes: 40 + device: h100 + num_devices: 4 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" + - label: MRCR Eval Small Models device: h200_35gb - timeout_in_minutes: 30 + timeout_in_minutes: 25 source_file_dependencies: - tests/evals/mrcr/ commands: diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 3ccf92f9a7a..46a3710ea80 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -5,18 +5,29 @@ steps: - label: LoRA %N device: h200_35gb key: lora - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/lora - tests/lora 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_qwen3_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 parallelism: 4 + mirror: + amd: + device: mi325_1 + working_dir: "/vllm-workspace/tests" + timeout_in_minutes: 65 + source_file_dependencies: + - vllm/lora + - tests/lora + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: LoRA TP (Distributed) key: lora-tp-distributed - timeout_in_minutes: 30 + timeout_in_minutes: 60 num_devices: 4 source_file_dependencies: - vllm/lora diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index e04016d6dcc..3365d09c7b8 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -5,7 +5,7 @@ steps: - label: V1 Spec Decode device: h200_35gb key: v1-spec-decode - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/config/ - vllm/distributed/ @@ -21,10 +21,16 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # TODO: create another `optional` test group for slow tests - pytest -v -s -m 'not slow_test' v1/spec_decode + mirror: + amd: + device: mi300_1 + timeout_in_minutes: 75 + depends_on: + - image-build-amd - label: V1 Sample + Logits key: v1-sample-logits - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/config/ @@ -58,7 +64,7 @@ steps: - label: V1 Core + KV + Metrics key: v1-core-kv-metrics - timeout_in_minutes: 30 + timeout_in_minutes: 60 source_file_dependencies: - vllm/config/ - vllm/distributed/ @@ -83,6 +89,7 @@ steps: - tests/v1/simple_kv_offload - tests/v1/worker - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py commands: @@ -95,10 +102,17 @@ steps: - pytest -v -s v1/simple_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/ec_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics # Integration test for streaming correctness (requires special branch). - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 75 + depends_on: + - image-build-amd - label: V1 Others (CPU) key: v1-others-cpu @@ -138,14 +152,29 @@ steps: - vllm/v1/spec_decode/extract_hidden_states.py - vllm/model_executor/models/extract_hidden_states.py - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py - tests/v1/kv_connector/extract_hidden_states_integration commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration +- label: Extract Hidden States Integration (2 GPUs) + key: extract-hidden-states-integration-2-gpus + timeout_in_minutes: 20 + num_devices: 2 + source_file_dependencies: + - vllm/v1/spec_decode/extract_hidden_states.py + - vllm/model_executor/models/extract_hidden_states.py + - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py + - tests/v1/kv_connector/extract_hidden_states_integration + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration + - label: Regression key: regression - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/config/ @@ -161,14 +190,14 @@ steps: - vllm/v1/ - tests/test_regression commands: - - pip install modelscope + - pip install 'modelscope<1.38' - pytest -v -s test_regression.py working_dir: "/vllm-workspace/tests" # optional - label: Examples device: h200_35gb key: examples - timeout_in_minutes: 45 + timeout_in_minutes: 40 working_dir: "/vllm-workspace/examples" source_file_dependencies: - vllm/entrypoints @@ -197,10 +226,20 @@ steps: - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048 # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536 + mirror: + amd: + device: mi325_1 + source_file_dependencies: + - vllm/entrypoints + - vllm/multimodal + - examples/ + - vllm/platforms/rocm.py + depends_on: + - image-build-amd - label: Metrics, Tracing (2 GPUs) key: metrics-tracing-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 25 num_devices: 2 source_file_dependencies: - vllm/config/ @@ -223,6 +262,12 @@ steps: 'opentelemetry-exporter-otlp>=1.26.0' \ 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing + mirror: + amd: + device: mi325_2 + depends_on: + - image-build-amd + optional: true - label: Python-only Installation key: python-only-installation @@ -235,11 +280,21 @@ steps: - setup.py commands: - bash standalone_tests/python_only_compile.sh + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - tests/standalone_tests/python_only_compile.sh + - setup.py + - vllm/platforms/rocm.py - label: Async Engine, Inputs, Utils, Worker device: h200_35gb key: async-engine-inputs-utils-worker - timeout_in_minutes: 50 + timeout_in_minutes: 25 source_file_dependencies: - vllm/assets/ - vllm/config/ @@ -266,7 +321,7 @@ steps: key: async-engine-inputs-utils-worker-config-cpu depends_on: - image-build-cpu - timeout_in_minutes: 30 + timeout_in_minutes: 65 source_file_dependencies: - vllm/assets/ - vllm/config/ @@ -293,38 +348,42 @@ steps: - vllm/transformers_utils/ - vllm/utils/ - vllm/v1/ + - tests/test_envs.py - tests/test_inputs.py - tests/test_outputs.py - tests/test_pooling_params.py - tests/test_ray_env.py + - tests/test_sampling_params.py - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py - - tests/tokenizers_ - tests/reasoning - tests/tool_parsers + - tests/tokenizers_ - tests/parser - tests/transformers_utils - tests/config device: cpu-small commands: - python3 standalone_tests/lazy_imports.py + - pytest -v -s test_envs.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 test_sampling_params.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 + - pytest -v -s reasoning - pytest -v -s tool_parsers + - pytest -v -s tokenizers_ - pytest -v -s parser - pytest -v -s transformers_utils - pytest -v -s config - label: Batch Invariance (A100) key: batch-invariance-a100 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: a100 source_file_dependencies: - vllm/v1/attention @@ -338,7 +397,7 @@ steps: - label: Batch Invariance (H100) key: batch-invariance-h100 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: h100 source_file_dependencies: - vllm/v1/attention @@ -354,7 +413,7 @@ steps: - label: Batch Invariance (B200) key: batch-invariance-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 35 device: b200-k8s source_file_dependencies: - vllm/v1/attention @@ -373,7 +432,7 @@ steps: - label: Acceptance Length Test (Large Models) # optional device: h200_35gb key: acceptance-length-test-large-models - timeout_in_minutes: 25 + timeout_in_minutes: 20 gpu: h100 optional: true num_gpus: 1 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index e34b7eadfac..4280e600df4 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Model Executor key: model-executor - timeout_in_minutes: 35 + timeout_in_minutes: 45 source_file_dependencies: - vllm/engine/arg_utils.py - vllm/config/model.py @@ -23,3 +23,16 @@ steps: # calls that the signal method cannot interrupt. - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread + mirror: + amd: + device: mi300_1 + depends_on: + - image-build-amd + 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 diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 617c80b2fec..3601aeee117 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -5,7 +5,7 @@ steps: - label: Model Runner V2 Core Tests device: h200_35gb key: model-runner-v2-core-tests - timeout_in_minutes: 45 + timeout_in_minutes: 35 source_file_dependencies: - vllm/v1/worker/gpu/ - vllm/v1/worker/gpu_worker.py @@ -18,9 +18,7 @@ steps: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" - # This requires eager until we sort out CG correctness issues. - # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged. - - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" + - pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" - pytest -v -s v1/e2e/general/test_context_length.py - pytest -v -s v1/e2e/general/test_min_tokens.py # Temporary hack filter to exclude ngram spec decoding based tests. @@ -29,7 +27,7 @@ steps: - label: Model Runner V2 Examples device: h200_35gb key: model-runner-v2-examples - timeout_in_minutes: 45 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/examples" source_file_dependencies: - vllm/v1/worker/gpu/ @@ -65,7 +63,7 @@ steps: - label: Model Runner V2 Distributed (2 GPUs) key: model-runner-v2-distributed-2-gpus - timeout_in_minutes: 45 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -86,7 +84,7 @@ steps: - label: Model Runner V2 Pipeline Parallelism (4 GPUs) key: model-runner-v2-pipeline-parallelism-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 4e47cbb7794..95827a89458 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -4,9 +4,8 @@ depends_on: steps: - label: Basic Models Tests (Initialization) key: basic-models-tests-initialization - timeout_in_minutes: 45 + timeout_in_minutes: 25 device: h200_18gb - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/test_initialization.py @@ -14,13 +13,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 device: h200_35gb key: basic-models-tests-extra-initialization - timeout_in_minutes: 45 + timeout_in_minutes: 100 source_file_dependencies: - vllm/model_executor/models/ - tests/models/test_initialization.py @@ -30,31 +27,35 @@ steps: # subset of supported models (the complement of the small subset in the above # 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: {} + parallelism: 4 - label: Basic Models Tests (Other) device: h200_35gb key: basic-models-tests-other - timeout_in_minutes: 45 + timeout_in_minutes: 35 source_file_dependencies: - vllm/ - tests/models/test_terratorch.py - - tests/models/test_transformers.py + - tests/models/transformers/test_backend.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py + mirror: + amd: + device: mi325_1 + depends_on: + - image-build-amd - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu depends_on: - image-build-cpu - timeout_in_minutes: 10 + timeout_in_minutes: 20 source_file_dependencies: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py + - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index b5758c55aff..c1ec5eb00ae 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed Model Tests (2 GPUs) key: distributed-model-tests-2-gpus - timeout_in_minutes: 50 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -17,7 +17,7 @@ steps: - 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)' # Avoid importing model tests that cause CUDA reinitialization error - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal/generation/test_phi4siglip.py -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 diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index f5e23cd95f4..d89d3fabaec 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Language Models Tests (Standard) key: language-models-tests-standard - timeout_in_minutes: 25 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/ @@ -14,11 +14,14 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' mirror: - torch_nightly: {} + amd: + device: mi300_1 + depends_on: + - image-build-amd - label: Language Models Tests (Extra Standard) %N key: language-models-tests-extra-standard - timeout_in_minutes: 45 + timeout_in_minutes: 40 source_file_dependencies: - vllm/model_executor/models/ - tests/models/language/pooling/test_embedding.py @@ -31,11 +34,25 @@ steps: - 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: {} + amd: + device: mi300_1 + depends_on: + - image-build-amd + source_file_dependencies: + - 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 - label: Language Models Tests (Hybrid) %N key: language-models-tests-hybrid - timeout_in_minutes: 75 + timeout_in_minutes: 65 source_file_dependencies: - vllm/ - tests/models/language/generation @@ -48,10 +65,9 @@ steps: - 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: {} amd: device: mi325_1 - timeout_in_minutes: 90 + timeout_in_minutes: 70 depends_on: - image-build-amd commands: @@ -62,7 +78,7 @@ steps: - label: Language Models Test (Extended Generation) # 80min device: h200_35gb key: language-models-test-extended-generation - timeout_in_minutes: 110 + timeout_in_minutes: 65 optional: true source_file_dependencies: - vllm/ @@ -76,7 +92,7 @@ steps: - label: Language Models Test (PPL) key: language-models-test-ppl - timeout_in_minutes: 110 + timeout_in_minutes: 30 device: h200_18gb optional: true source_file_dependencies: @@ -88,7 +104,7 @@ steps: - label: Language Models Test (Extended Pooling) # 36min device: h200_35gb key: language-models-test-extended-pooling - timeout_in_minutes: 50 + timeout_in_minutes: 70 optional: true source_file_dependencies: - vllm/ @@ -104,7 +120,7 @@ steps: - label: Language Models Test (MTEB) key: language-models-test-mteb - timeout_in_minutes: 110 + timeout_in_minutes: 45 device: h200_18gb optional: true source_file_dependencies: diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 7c7e2163fef..473192f03f9 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -10,93 +10,94 @@ steps: - 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 mirror: amd: - device: mi300_1 + device: mi325_1 depends_on: - image-build-amd - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" key: multi-modal-models-standard-2-qwen3-gemma - timeout_in_minutes: 45 + timeout_in_minutes: 50 device: h200_18gb 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_mm_prefix_lm.py -m core_model - 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: mi300_1 + device: mi325_1 depends_on: - image-build-amd - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" device: h200_35gb key: multi-modal-models-standard-3-llava-qwen2-vl - timeout_in_minutes: 45 + timeout_in_minutes: 40 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 mirror: amd: - device: mi300_1 + device: mi325_1 depends_on: - image-build-amd - label: "Multi-Modal Models (Standard) 4: other + whisper" device: h200_35gb key: multi-modal-models-standard-4-other-whisper - timeout_in_minutes: 45 + timeout_in_minutes: 50 source_file_dependencies: - vllm/ - tests/models/multimodal commands: - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -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_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - 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: + device: mi325_1 + depends_on: + - image-build-amd -- label: Multi-Modal Processor (CPU) +- label: Multi-Modal Processor (CPU) %N key: multi-modal-processor-cpu depends_on: - image-build-cpu - timeout_in_minutes: 60 + timeout_in_minutes: 125 source_file_dependencies: - vllm/ - tests/models/multimodal - tests/models/registry.py device: cpu-medium 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 + - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + parallelism: 4 - label: Multi-Modal Processor # 44min key: multi-modal-processor - timeout_in_minutes: 60 + timeout_in_minutes: 65 device: h200_18gb 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 Accuracy Eval (Small Models) # 50min device: h200_35gb key: multi-modal-accuracy-eval-small-models - timeout_in_minutes: 70 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - vllm/multimodal/ @@ -104,6 +105,17 @@ steps: - vllm/v1/core/ commands: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 + mirror: + amd: + device: mi300_1 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/multimodal/ + - vllm/inputs/ + - vllm/v1/core/ + - vllm/platforms/rocm.py + - vllm/model_executor/model_loader/ - label: Multi-Modal Models (Extended Generation 1) key: multi-modal-models-extended-generation-1 @@ -113,7 +125,6 @@ steps: - 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 mirror: @@ -130,7 +141,6 @@ steps: - 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) @@ -141,7 +151,6 @@ steps: - 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 (Extended Pooling) @@ -155,8 +164,8 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' mirror: amd: - device: mi300_1 - timeout_in_minutes: 60 + device: mi325_1 + timeout_in_minutes: 75 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 591afd946d2..e0bb67b1bb4 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Plugin Tests (2 GPUs) key: plugin-tests-2-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -27,12 +27,21 @@ steps: - 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 + # test colbert_query io_processor plugin + - pip install -e ./plugins/colbert_query_plugin + - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py + - pip uninstall colbert_query_plugin -y # end io_processor plugins 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 endpoint plugins test + - pip install -e ./plugins/vllm_add_dummy_endpoint_plugin + - pytest -v -s plugins_tests/test_endpoint_plugins.py + - pip uninstall vllm_add_dummy_endpoint_plugin -y + # end endpoint plugins test # other tests continue here: - pytest -v -s plugins_tests/test_scheduler_plugins.py - pip install -e ./plugins/vllm_add_dummy_model @@ -40,3 +49,17 @@ steps: - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins + + +- label: GGUF Plugin + key: gguf-plugin + device: h200_18gb + timeout_in_minutes: 30 + soft_fail: true + optional: true + source_file_dependencies: + - vllm/model_executor/layers/quantization + - tests/plugins_tests/test_gguf_plugin.py + commands: + - pip install "vllm-gguf-plugin >= 0.0.2" + - pytest -v -s plugins_tests/gguf diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 6866d5e3695..72cb0069664 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -5,7 +5,7 @@ steps: - label: PyTorch Compilation Unit Tests device: h200_35gb key: pytorch-compilation-unit-tests - timeout_in_minutes: 10 + timeout_in_minutes: 90 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -78,7 +78,7 @@ steps: - label: PyTorch Compilation Passes Unit Tests key: pytorch-compilation-passes-unit-tests - timeout_in_minutes: 20 + timeout_in_minutes: 45 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -107,10 +107,16 @@ steps: - tests/compile/passes commands: - pytest -s -v compile/passes --ignore compile/passes/distributed + mirror: + amd: + device: mi300_1 + timeout_in_minutes: 65 + depends_on: + - image-build-amd - label: PyTorch Fullgraph Smoke Test key: pytorch-fullgraph-smoke-test - timeout_in_minutes: 35 + timeout_in_minutes: 60 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -146,7 +152,7 @@ steps: - label: PyTorch Fullgraph key: pytorch-fullgraph - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: h200_18gb source_file_dependencies: - vllm/__init__.py @@ -189,3 +195,11 @@ steps: - requirements/test/nightly-torch.txt commands: - bash standalone_tests/pytorch_nightly_dependency.sh + mirror: + amd: + device: mi300_1 + depends_on: + - image-build-amd + source_file_dependencies: + - requirements/test/nightly-torch.txt + - vllm/platforms/rocm.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 8a9a36da448..ce3e58e501b 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Quantization key: quantization - timeout_in_minutes: 90 + timeout_in_minutes: 60 source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -21,9 +21,21 @@ steps: - 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 Fusions + key: quantized-fusions + timeout_in_minutes: 20 + source_file_dependencies: + - tests/fusion + - vllm/model_executor/layers/fusion + - vllm/model_executor/kernels/linear + - vllm/model_executor/layers/quantization/compressed_tensors + - vllm/model_executor/layers/quantization/modelopt.py + commands: + - pytest -v -s fusion/ + - label: Quantized MoE Test (B200) key: quantized-moe-test-b200 - timeout_in_minutes: 60 + timeout_in_minutes: 120 working_dir: "/vllm-workspace/" device: b200-k8s source_file_dependencies: @@ -41,7 +53,7 @@ steps: - label: Quantized Models Test key: quantized-models-test - timeout_in_minutes: 60 + timeout_in_minutes: 50 source_file_dependencies: - vllm/model_executor/layers/quantization - tests/models/quantization diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 16d69f77345..9e5e09c3ec3 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -3,7 +3,7 @@ depends_on: - image-build steps: - label: Rust Frontend OpenAI Coverage - timeout_in_minutes: 90 + timeout_in_minutes: 30 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -15,28 +15,30 @@ steps: - tests/utils.py - tests/benchmarks/test_serve_cli.py - tests/entrypoints/openai/chat_completion/test_chat_completion.py - # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py + - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py # - tests/entrypoints/openai/completion/test_prompt_validation.py - tests/entrypoints/openai/completion/test_shutdown.py - # - tests/entrypoints/openai/test_return_token_ids.py - # - tests/entrypoints/openai/test_uds.py + - tests/entrypoints/openai/test_return_token_ids.py + - tests/entrypoints/openai/test_uds.py - tests/v1/sample/test_logprobs_e2e.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py - # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds" - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" - # - pytest -v -s entrypoints/openai/test_return_token_ids.py - # - pytest -v -s entrypoints/openai/test_uds.py + # test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and + # finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk. + - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" + - pytest -v -s entrypoints/openai/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" - label: Rust Frontend Serve/Admin Coverage - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -45,22 +47,27 @@ steps: - vllm/entrypoints/serve/ - vllm/v1/engine/ - tests/utils.py - # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py - - tests/entrypoints/serve/disagg/test_serving_tokens.py + - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py + - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py # - tests/entrypoints/serve/dev/test_sleep.py + - tests/entrypoints/serve/tokenize/test_tokenization.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + # server_load can be flaky under the Rust frontend; keep it excluded for now. - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" - - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" + # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). + - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" # - pytest -v -s entrypoints/serve/dev/test_sleep.py + # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" - label: Rust Frontend Core Correctness - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -74,7 +81,7 @@ steps: - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - label: Rust Frontend Tool Use - timeout_in_minutes: 60 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" source_file_dependencies: - rust/ @@ -88,7 +95,7 @@ steps: - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" - label: Rust Frontend Distributed - timeout_in_minutes: 30 + timeout_in_minutes: 25 num_devices: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -99,9 +106,13 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/utils.py + - tests/v1/distributed/test_external_lb_dp.py + - tests/v1/distributed/test_hybrid_lb_dp.py - tests/v1/distributed/test_internal_lb_dp.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export NCCL_CUMEM_HOST_ENABLE=0 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 06f9eb9c245..21d4c2ac219 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -4,7 +4,7 @@ steps: - label: Rust Frontend Cargo Style + Clippy key: rust-frontend-cargo-style-clippy depends_on: [] - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: cpu-medium no_plugin: true source_file_dependencies: @@ -18,7 +18,7 @@ steps: - label: Rust Frontend Cargo Tests key: rust-frontend-cargo-tests depends_on: [] - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: cpu-medium no_plugin: true source_file_dependencies: diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 6ec6f8efd35..2e7cd4a623e 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -5,7 +5,7 @@ steps: - label: Samplers Test device: h200_35gb key: samplers-test - timeout_in_minutes: 75 + timeout_in_minutes: 40 source_file_dependencies: - vllm/model_executor/layers - vllm/sampling_metadata.py @@ -19,7 +19,7 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers mirror: amd: - device: mi250_1 + device: mi325_1 depends_on: - image-build-amd commands: diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 27396118dbc..096c324bb8e 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Spec Decode Eagle key: spec-decode-eagle - timeout_in_minutes: 30 + timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -12,10 +12,24 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness" + mirror: + amd: + device: mi325_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + 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 - label: Spec Decode Eagle Nightly B200 key: spec-decode-eagle-nightly-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 25 device: b200-k8s optional: true source_file_dependencies: @@ -27,7 +41,7 @@ steps: - label: Spec Decode Speculators + MTP key: spec-decode-speculators-mtp - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -39,7 +53,7 @@ steps: - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" mirror: amd: - device: mi300_1 + device: mi325_1 timeout_in_minutes: 65 depends_on: - image-build-amd @@ -68,7 +82,7 @@ steps: - label: Spec Decode Ngram + Suffix key: spec-decode-ngram-suffix - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -78,8 +92,10 @@ steps: - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" mirror: amd: - device: mi300_1 - timeout_in_minutes: 65 + device: mi325_1 + timeout_in_minutes: 55 + # TODO(akaratza): Test after Torch >= 2.12 bump + soft_fail: true depends_on: - image-build-amd source_file_dependencies: @@ -93,7 +109,7 @@ steps: - label: Spec Decode Draft Model key: spec-decode-draft-model - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -103,8 +119,8 @@ steps: - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" mirror: amd: - device: mi300_1 - timeout_in_minutes: 50 + device: mi325_1 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: @@ -118,7 +134,7 @@ steps: - label: Spec Decode Draft Model Nightly B200 key: spec-decode-draft-model-nightly-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: b200-k8s optional: true source_file_dependencies: @@ -130,7 +146,7 @@ steps: - label: Speculators Correctness key: speculators-correctness - timeout_in_minutes: 60 + timeout_in_minutes: 30 device: h100 optional: true num_devices: 1 @@ -143,7 +159,7 @@ steps: - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test - label: Spec Decode MTP hybrid (B200) - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: b200-k8s optional: true source_file_dependencies: diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 01c6bb7809b..eeb24a49d8d 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Weight Loading Multiple GPU # 33min key: weight-loading-multiple-gpu - timeout_in_minutes: 45 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" num_devices: 2 optional: true @@ -13,6 +13,13 @@ steps: - tests/weight_loading commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt + mirror: + amd: + device: mi300_2 + depends_on: + - image-build-amd + commands: + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt # - label: Weight Loading Multiple GPU - Large Models # optional # working_dir: "/vllm-workspace/tests" diff --git a/.claude/skills/ci-fails-buildkite/SKILL.md b/.claude/skills/ci-fails-buildkite/SKILL.md new file mode 100644 index 00000000000..d195c02f723 --- /dev/null +++ b/.claude/skills/ci-fails-buildkite/SKILL.md @@ -0,0 +1,35 @@ +--- +name: ci-fails-buildkite +description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs. +--- + +# Diagnosing vLLM Buildkite CI Failures + +Buildkite logs are public; no login needed. + +`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci--.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch. + +## Fetching logs + +```bash +# All failed jobs in a PR's latest build (current branch's PR if omitted): +.buildkite/scripts/ci-fetch-log.sh --pr + +# All failed jobs in a build (--soft also includes soft-failed jobs; +# --all fetches every finished job): +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/" + +# One job — `gh pr checks` URLs (#) and web UI URLs (?sid=) both +# work; pass "-" as a second argument to stream to stdout: +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/#" +``` + +To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`: + +```bash +./ci-clean-log.sh ci.log +``` + +## Reference + +See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage. diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a8947fe2324..57166d9d9b7 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,17 +2,16 @@ # for more info about CODEOWNERS file # This lists cover the "core" components of vLLM that require careful review -/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy -/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi +/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng +/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye /vllm/model_executor/layers/mamba @tdoublep @tomeras91 -/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy -/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy +/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn -/vllm/model_executor/layers/batch_invariant.py @yewentao256 +/vllm/model_executor/layers/batch_invariant.py @yewentao256 /vllm/ir @ProExpertProg /vllm/kernels/ @ProExpertProg @tjtanaa /vllm/kernels/helion @ProExpertProg @zou3519 @@ -23,8 +22,13 @@ # Any change to the VllmConfig changes can have a large user-facing impact, # so spam a lot of people -/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg -/vllm/config/cache.py @heheda12345 +/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg +/vllm/config/cache.py @heheda12345 @ivanium + +# Config utils +/vllm/config/utils.py @hmellor +/vllm/engine/arg_utils.py @hmellor +/vllm/utils/argparse_utils.py # Entrypoints /vllm/entrypoints/anthropic @mgoin @DarkLight1337 @@ -63,19 +67,20 @@ /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety @vadiklyutiy /vllm/v1/attention/backends/triton_attn.py @tdoublep /vllm/v1/attention/backends/gdn_attn.py @ZJY0516 @vadiklyutiy -/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery +/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett -/vllm/v1/kv_cache_interface.py @heheda12345 +/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery +/vllm/v1/simple_kv_offload @ivanium /vllm/v1/engine @njhill /vllm/v1/executor @njhill /vllm/v1/worker @njhill -/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche +/vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium # Model runner V2 -/vllm/v1/worker/gpu @WoosukKwon @njhill +/vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256 /vllm/v1/worker/gpu/kv_connector.py @orozery # CI & building @@ -99,13 +104,14 @@ /tests/test_inputs.py @DarkLight1337 @ywang96 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm /tests/v1/structured_output @mgoin @russellb @aarnphm -/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery +/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee /tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 /tests/v1/kv_connector/nixl_integration @NickLucche -/tests/v1/kv_connector @ApostaC @orozery +/tests/v1/kv_connector @ApostaC @orozery @ivanium /tests/v1/kv_offload @ApostaC @orozery +/tests/v1/simple_kv_offload @ivanium /tests/v1/determinism @yewentao256 /tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning /tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning @@ -113,17 +119,7 @@ # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor -/tests/models/test_transformers.py @hmellor - -# Observability -/vllm/config/observability.py @markmc -/vllm/v1/metrics @markmc -/tests/v1/metrics @markmc -/vllm/tracing.py @markmc -/tests/v1/tracing/test_tracing.py @markmc -/vllm/config/kv_events.py @markmc -/vllm/distributed/kv_events.py @markmc -/tests/distributed/test_events.py @markmc +/tests/models/transformers @hmellor # Docs /docs/mkdocs @hmellor diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml new file mode 100644 index 00000000000..082e8a9eb90 --- /dev/null +++ b/.github/actionlint.yaml @@ -0,0 +1,7 @@ +# Custom self-hosted runner labels (e.g. the autoscaling vllm-runners pool) so +# actionlint doesn't flag them as unknown in `runs-on`. +self-hosted-runner: + labels: + - vllm-runners + # Not yet in actionlint's known-label set. + - macos-26 diff --git a/.github/dependabot.yml b/.github/dependabot.yml index a017d69be99..944929fc55e 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -21,7 +21,6 @@ updates: - dependency-name: "torchvision" - dependency-name: "xformers" - dependency-name: "lm-format-enforcer" - - dependency-name: "gguf" - dependency-name: "compressed-tensors" - dependency-name: "ray[cgraph]" # Ray Compiled Graph - dependency-name: "lm-eval" diff --git a/.github/mergify.yml b/.github/mergify.yml index e245cf6baca..4333c6e646d 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -144,12 +144,12 @@ pull_request_rules: - label != stale - or: - files~=^examples/.*mistral.*\.py - - files~=^tests/.*mistral.*\.py - - files~=^vllm/model_executor/models/.*mistral.*\.py + - files~=^tests/.*(?:mistral|voxtral|mixtral|pixtral).*\.py + - files~=^vllm/model_executor/models/.*(?:mistral|voxtral|mixtral|pixtral).*\.py - files~=^vllm/reasoning/.*mistral.*\.py - files~=^vllm/tool_parsers/.*mistral.*\.py - - files~=^vllm/transformers_utils/.*mistral.*\.py - - title~=(?i)Mistral + - files~=^vllm/transformers_utils/.*(?:mistral|voxtral|pixtral).*\.py + - title~=(?i)(?:mistral|ministral|voxtral|mixtral|pixtral) actions: label: add: @@ -388,9 +388,13 @@ pull_request_rules: - or: - files~=^tests/tool_use/ - files~=^tests/tool_parsers/ + - files~=^tests/parser/ + - files~=^tests/reasoning/ - files~=^tests/entrypoints/openai/.*tool.* - files~=^tests/entrypoints/anthropic/.*tool.* - files~=^vllm/tool_parsers/ + - files~=^vllm/parser/ + - files~=^vllm/reasoning/ - files=docs/features/tool_calling.md - files~=^examples/tool_calling/ actions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 4eac3d7b789..7a98ce7cc08 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -327,7 +327,7 @@ jobs: message: 'CC {users} for ROCm-related issue', }, mistral: { - users: ['patrickvonplaten', 'juliendenize', 'andylolu2'], + users: ['patrickvonplaten', 'juliendenize', 'andylolu2', 'NickLucche'], message: 'CC {users} for Mistral-related issue', }, // Add more label -> user mappings here diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index ea1c8b0feac..011bf84feb3 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -11,13 +11,25 @@ permissions: jobs: macos-m1-smoke-test: - runs-on: macos-latest + # macos-26 (the supported target) is still a preview runner, so gate on GA + # macos-15 and keep macos-26 non-blocking. + strategy: + fail-fast: false + matrix: + include: + - os: macos-15 + required: true + - os: macos-26 + required: false + name: macos-m1-smoke-test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + continue-on-error: ${{ !matrix.required }} timeout-minutes: 30 steps: - - uses: actions/checkout@v6.0.1 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: astral-sh/setup-uv@v7 + - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: enable-cache: true cache-dependency-glob: | @@ -72,14 +84,11 @@ jobs: # Test health endpoint curl -f http://localhost:8000/health - # Test completion - curl -f http://localhost:8000/v1/completions \ + # Long prompt: hits the split-KV path that short prompts skip (#46769). + PAYLOAD=$(python -c "import json; print(json.dumps({'model': 'Qwen/Qwen3-0.6B', 'prompt': 'The quick brown fox jumps over the lazy dog. ' * 24, 'max_tokens': 16}))") + curl -f --max-time 120 http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ - -d '{ - "model": "Qwen/Qwen3-0.6B", - "prompt": "Hello", - "max_tokens": 5 - }' + -d "$PAYLOAD" # Cleanup kill "$SERVER_PID" diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 93a5a5ff0ae..143fc427a49 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -28,7 +28,8 @@ jobs: pull_number: context.payload.pull_request.number, }); - const hasReadyLabel = pr.labels.some(l => l.name === 'ready'); + const readyLabels = ['ready', 'ready-run-all-tests']; + const hasReadyLabel = pr.labels.some(l => readyLabels.includes(l.name)); const hasVerifiedLabel = pr.labels.some(l => l.name === 'verified'); const { data: mergedPRs } = await github.rest.search.issuesAndPullRequests({ @@ -40,18 +41,22 @@ jobs: if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) { core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`); } else { - core.setFailed(`PR must have the 'verified' or 'ready' (which also triggers tests) label or the author must have at least 4 merged PRs (found ${mergedCount}).`); + core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`); } pre-commit: needs: pre-run-check if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') - runs-on: ubuntu-latest + runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 with: python-version: "3.12" + # Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its + # wget + tar -xJ self-download, which the self-hosted runner image lacks + # (no wget/xz). Pinned to shellcheck 0.10.0 to match the script's "stable". + - run: python -m pip install shellcheck-py==0.10.0.1 - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json" - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" diff --git a/.github/workflows/scripts/build.sh b/.github/workflows/scripts/build.sh index eb3971c42bf..335ec735e62 100644 --- a/.github/workflows/scripts/build.sh +++ b/.github/workflows/scripts/build.sh @@ -9,7 +9,7 @@ PATH=${cuda_home}/bin:$PATH LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH # Install requirements -if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then +if [ "$(echo "$2" | cut -d. -f1)" = "12" ]; then sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt fi $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt @@ -17,7 +17,10 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements # Limit the number of parallel jobs to avoid OOM export MAX_JOBS=1 # Make sure release wheels are built for the following architectures -export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" +# Do not add +PTX here: vLLM filters torch's top-level PTX flag when it +# converts global gencode flags into per-kernel arch lists. If a specific +# kernel needs PTX, add +PTX to that kernel's CMake arch list instead. +export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0" bash tools/check_repo.sh diff --git a/.gitignore b/.gitignore index 2c4e135e58d..26cd21a015d 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py # DeepGEMM vendored package built from source vllm/third_party/deep_gemm/ +# fmha_sm100 vendored package built from source +vllm/third_party/fmha_sm100/ + # triton jit .triton @@ -196,7 +199,9 @@ cython_debug/ .vscode/ # Claude -.claude/ +.claude/* +!.claude/skills/ +!.claude/skills/** # Codex .codex/ @@ -233,7 +238,7 @@ actionlint shellcheck*/ # Ignore moe/marlin_moe gen code -csrc/moe/marlin_moe_wna16/kernel_* +csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_* # Ignore ep_kernels_workspace folder ep_kernels_workspace/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dff099e3697..ad767c823be 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: rev: v21.1.2 hooks: - id: clang-format - exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*' + exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 @@ -131,6 +131,19 @@ repos: --python-version, "3.12", ] files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$ + - id: pip-compile + alias: pip-compile-cpu + name: pip-compile-cpu + args: [ + requirements/test/cuda.in, + -o, requirements/test/cpu.txt, + --index-strategy, unsafe-best-match, + --torch-backend, cpu, + --python-platform, x86_64-manylinux_2_28, + --python-version, "3.12", + ] + files: ^requirements/(common|cpu|test/(cuda|cpu))\.(in|txt)$ + exclude: ^requirements/test/cuda\.txt$ - id: pip-compile alias: pip-compile-docs name: pip-compile-docs @@ -197,7 +210,7 @@ repos: name: Check SPDX headers entry: python tools/pre_commit/check_spdx_header.py language: python - types: [python] + types_or: [python, rust, proto] - id: check-root-lazy-imports name: Check root lazy imports entry: python tools/pre_commit/check_init_lazy_imports.py diff --git a/AGENTS.md b/AGENTS.md index 441b8d9fb73..a53b81873cf 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -29,6 +29,7 @@ Do not open one-off PRs for tiny edits (single typo, isolated style change, one - PR descriptions for AI-assisted work **must** include: - Why this is not duplicating an existing PR. - Test commands run and results. + - Model evaluation results when the change affects output, accuracy, or serving. - Clear statement that AI assistance was used. ### Fail-closed behavior @@ -66,23 +67,38 @@ VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto uv pip install -e . --torch-backend=auto ``` -### Running tests +### Tests > Requires [Environment setup](#environment-setup) and [Installing dependencies](#installing-dependencies). ```bash -# Install test dependencies. -# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the -# unpinned source file instead: -uv pip install -r requirements/test/cuda.in # resolves for current platform -# Or on x86_64: -uv pip install -r requirements/test/cuda.txt +# Install test dependencies (use cuda.in on non-x86_64): +uv pip install -r requirements/test/cuda.in -# Run a specific test file (use .venv/bin/python directly; -# `source activate` does not persist in non-interactive shells): +# Run a specific test file: .venv/bin/python -m pytest tests/path/to/test_file.py -v ``` +When adding tests: + +- **Design before you write.** Answer four questions first: what is the module + for, what is its I/O contract, what failure am I guarding against, and what is + the cheapest level that catches it (unit over integration over e2e)? +- **Reuse before create.** Extend existing test files, `conftest.py` fixtures, and + helpers; add a new file only when no nearby suite fits. +- **Test behavior with intent.** Assert observable outcomes through public APIs; + state why in the name or docstring. Skip trivial wiring; flaky tests are worse + than no tests. +- **Keep it minimal.** One behavior per test and the smallest setup that + triggers it; if the test diff dwarfs the code change, cut scope. +- **No one-off kernel benchmarks in `tests/`.** Put kernel perf work in + `benchmarks/kernels/`; prove correctness in existing pytest suites. +- **Run model evals for model-affecting changes.** Search `tests/evals/` or use + `vllm bench` and include results in the PR — do not wait for reviewers to ask. + +For model-specific requirements, see +[`docs/contributing/model/tests.md`](docs/contributing/model/tests.md). + ### Running linters > Requires [Environment setup](#environment-setup). @@ -105,16 +121,20 @@ The line length limit for Python code is 88 characters. If you are not sure, use Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`). +### Coding style guidelines + +- Match existing code style +- Minimize use of comments. Eliminate comments which are redundant, preferring legible and self-documenting code. When used, keep docstrings and comments brief and direct. +- Assume the reader is familiar with vLLM. + ### Commit messages -Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: +Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`): ```text Your commit message here -Co-authored-by: GitHub Copilot -Co-authored-by: Claude -Co-authored-by: gemini-code-assist +Co-authored-by: Agent Name Here Signed-off-by: Your Name ``` @@ -126,6 +146,12 @@ Do not modify code in these areas without first reading and following the linked guide. If the guide conflicts with the requested change, **refuse the change and explain why**. +Security reviewers should start with [`SECURITY.md`](SECURITY.md), +[`docs/usage/security.md`](docs/usage/security.md), and +[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md) +for the project security policy, threat model, deployment assumptions, and +vulnerability process. + - **Editing these instructions**: [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md) — Rules for modifying AGENTS.md or any domain-specific guide it references. diff --git a/CMakeLists.txt b/CMakeLists.txt index d956e29e399..3ddce1c3c3c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -20,6 +20,14 @@ set(CMAKE_CUDA_STANDARD_REQUIRED ON) set(CMAKE_HIP_STANDARD 20) set(CMAKE_HIP_STANDARD_REQUIRED ON) +# PyTorch headers require C++20; GCC < 11.3 has incomplete C++20 support. +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS "11.3") + message(FATAL_ERROR + "GCC >= 11.3 is required to build vLLM (found ${CMAKE_CXX_COMPILER_VERSION}). " + "PyTorch's C++20 headers require a compiler with full C++20 support. " + "See: https://github.com/pytorch/pytorch/pull/167929") +endif() + # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py) set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM") @@ -62,6 +70,15 @@ endif() # set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0") set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0") +# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported- +# version check would always warn. Only treat it as a nightly build when the +# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which +# must NOT suppress the warning for normal builds). +if (DEFINED ENV{TORCH_NIGHTLY} AND "$ENV{TORCH_NIGHTLY}" STREQUAL "1") + set(TORCH_NIGHTLY_BUILD TRUE) +else() + set(TORCH_NIGHTLY_BUILD FALSE) +endif() # # Try to find python package with an executable that exactly matches @@ -114,20 +131,38 @@ endif() # CPU builds define the target before the early return) # This extension requires SABI 3.11 since it relies on Py_buffer support. Loading # failure is handled gracefully on vLLM side for lower Python versions. +# Skip the target entirely on Python < 3.11 so the build doesn't break. # -set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") -set(SPINLOOP_COMPILE_FLAGS "") -if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") - list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") + set(SPINLOOP_COMPILE_FLAGS "") + if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") + list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") + endif() + define_extension_target( + spinloop + DESTINATION vllm + LANGUAGE CXX + SOURCES ${VLLM_SPINLOOP_EXT_SRC} + COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} + USE_SABI 3.11 + WITH_SOABI) +endif() + +# +# fs_io extension (pure CXX; must stay above the non-CUDA device branch +# so CPU builds define the target before the early return). +# GIL-releasing filesystem helpers for FileSystemTierManager. +# +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + define_extension_target( + fs_io_C + DESTINATION vllm + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) endif() -define_extension_target( - spinloop - DESTINATION vllm - LANGUAGE CXX - SOURCES ${VLLM_SPINLOOP_EXT_SRC} - COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} - USE_SABI 3.11 - WITH_SOABI) # # Forward the non-CUDA device extensions to external CMake scripts. @@ -149,7 +184,7 @@ endif() if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND) set(VLLM_GPU_LANG "CUDA") - if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) + if (NOT TORCH_NIGHTLY_BUILD AND NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} " "expected for CUDA build, saw ${Torch_VERSION} instead.") endif() @@ -162,7 +197,7 @@ elseif(HIP_FOUND OR PYTORCH_FOUND_HIP) enable_language(HIP) # ROCm 5.X and 6.X - if (ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND + if (NOT TORCH_NIGHTLY_BUILD AND ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM}) message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} " "expected for ROCm build, saw ${Torch_VERSION} instead.") @@ -179,6 +214,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # the set of architectures we want to compile for and remove the from the # CMAKE_CUDA_FLAGS so that they are not applied globally. # + # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch + # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only + # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's + # component-specific arch list below. + # clear_cuda_arches(CUDA_ARCH_FLAGS) extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") message(STATUS "CUDA target architectures: ${CUDA_ARCHS}") @@ -254,6 +294,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP") # 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") + + # When using LTO then *.cpp files must be compiled with same compiler as used linker + # So if HIP uses clang linker we also must use it + # Otherwise symbols will be missing from .so + if (CMAKE_CXX_FLAGS MATCHES "\-flto") + if(NOT CMAKE_CXX_COMPILER_ID STREQUAL CMAKE_HIP_COMPILER_ID) + message(FATAL_ERROR "LTO is enabled for ROCm build, but the C++ compiler (${CMAKE_CXX_COMPILER_ID}) and HIP compiler (${CMAKE_HIP_COMPILER_ID}) are different which is not supported. " + "Please ensure they are same by setting CXX=${CMAKE_HIP_COMPILER} environment variable. Or alternatively disable LTO.") + endif() + endif() endif() # @@ -303,320 +353,33 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # -# _C extension +# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch) # -set(VLLM_EXT_SRC - "csrc/quantization/activation_kernels.cu" - "csrc/torch_bindings.cpp") +if(VLLM_GPU_LANG STREQUAL "HIP") + set(VLLM_EXT_SRC + "csrc/torch_bindings.cpp" + "csrc/custom_quickreduce.cu") -if(VLLM_GPU_LANG STREQUAL "CUDA") - SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") + message(STATUS "Enabling C extension.") + define_extension_target( + _C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${VLLM_EXT_SRC} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} + INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} + USE_SABI 3 + WITH_SOABI) - # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. - set(CUTLASS_REVISION "v4.4.2") - - # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided - if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR}) - set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR}) - endif() - - if(VLLM_CUTLASS_SRC_DIR) - if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR) - get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE) - endif() - message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation") - FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR}) - else() - FetchContent_Declare( - cutlass - GIT_REPOSITORY https://github.com/nvidia/cutlass.git - # Please keep this in sync with CUTLASS_REVISION line above. - GIT_TAG ${CUTLASS_REVISION} - GIT_PROGRESS TRUE - - # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. - # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. - # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE - GIT_SHALLOW TRUE - ) - endif() - FetchContent_MakeAvailable(cutlass) - - set_gencode_flags_for_srcs( - SRCS "${VLLM_EXT_SRC}" - CUDA_ARCHS "${CUDA_ARCHS}") - - # Only build Marlin kernels if we are building for at least some compatible archs. - # Keep building Marlin for 9.0 as there are some group sizes and shapes that - # are not supported by Machete yet. - - # marlin arches for fp16 output - # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0; - # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8). - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin has limited support for turing - cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}") - # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX) - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin arches for fp8 input - # - sm80 doesn't support fp8 computation - # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction - # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin arches for other files - cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") - - if (MARLIN_OTHER_ARCHS) - - # - # For the Marlin kernels we automatically generate sources for various - # preselected input type pairs and schedules. - # Generate sources: - set(MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) - file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) - list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) - set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") - - message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") - message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") - - if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} - OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}) - execute_process( - COMMAND ${CMAKE_COMMAND} -E env - PYTHONPATH=$ENV{PYTHONPATH} - ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} - RESULT_VARIABLE marlin_generation_result - OUTPUT_VARIABLE marlin_generation_result - OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log - ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log - ) - - if (NOT marlin_generation_result EQUAL 0) - message(FATAL_ERROR "Marlin generation failed." - " Result: \"${marlin_generation_result}\"" - "\nCheck the log for details: " - "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log") - else() - set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} - CACHE STRING "Last run Marlin generate script hash and arch" FORCE) - message(STATUS "Marlin generation completed successfully.") - endif() - else() - message(STATUS "Marlin generation script has not changed, skipping generation.") - endif() - - if (MARLIN_ARCHS) - file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) - - file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_BF16_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC}) - endif() - - if (MARLIN_SM75_ARCHS) - file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_SM75_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC}) - endif() - - if (MARLIN_FP8_ARCHS) - file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_FP8_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC}) - endif() - - set(MARLIN_SRCS - "csrc/quantization/marlin/marlin.cu" - "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu" - "csrc/quantization/marlin/gptq_marlin_repack.cu" - "csrc/quantization/marlin/awq_marlin_repack.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_SRCS}" - CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_SRCS} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}") - - message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}") - else() - message(STATUS "Not building Marlin kernels as no compatible archs found" - " in CUDA target architectures") - endif() - - # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+). - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") - endif() - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS) - set(SRCS - "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu" - "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu") - set_gencode_flags_for_srcs( - SRCS "${SRCS}" - CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1") - message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}") - else() - if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 - AND ES_MXFP8_GROUPED_MM_ARCHS) - message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is " - "not >= 12.8.") - else() - message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found " - "in CUDA target architectures.") - endif() - endif() - - # - # Machete kernels - - # The machete kernels only work on hopper and require CUDA 12.0 or later. - # Only build Machete kernels if we are building for something compatible with sm90a - cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS) - # - # For the Machete kernels we automatically generate sources for various - # preselected input type pairs and schedules. - # Generate sources: - set(MACHETE_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py) - file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) - - message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") - message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}") - - if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH} - OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH}) - execute_process( - COMMAND ${CMAKE_COMMAND} -E env - PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH} - ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT} - RESULT_VARIABLE machete_generation_result - OUTPUT_VARIABLE machete_generation_output - OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log - ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log - ) - - if (NOT machete_generation_result EQUAL 0) - message(FATAL_ERROR "Machete generation failed." - " Result: \"${machete_generation_result}\"" - "\nCheck the log for details: " - "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log") - else() - set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH} - CACHE STRING "Last run machete generate script hash" FORCE) - message(STATUS "Machete generation completed successfully.") - endif() - else() - message(STATUS "Machete generation script has not changed, skipping generation.") - endif() - - # Add machete generated sources - file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu") - list(APPEND VLLM_EXT_SRC ${MACHETE_GEN_SOURCES}) - - # forward compatible - set_gencode_flags_for_srcs( - SRCS "${MACHETE_GEN_SOURCES}" - CUDA_ARCHS "${MACHETE_ARCHS}") - - list(APPEND VLLM_EXT_SRC - csrc/quantization/machete/machete_pytorch.cu) - - message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") - else() - if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 - AND MACHETE_ARCHS) - message(STATUS "Not building Machete kernels as CUDA Compiler version is " - "not >= 12.0, we recommend upgrading to CUDA 12.0 or " - "later if you intend on running w4a16 quantized models on " - "Hopper.") - else() - message(STATUS "Not building Machete kernels as no compatible archs " - "found in CUDA target architectures") - endif() - endif() - - - -# if CUDA endif -endif() - -if (VLLM_GPU_LANG STREQUAL "HIP") - # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration). - # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11. - list(APPEND VLLM_EXT_SRC - "csrc/custom_quickreduce.cu" - "csrc/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" - ) -# if ROCM endif -endif() - -message(STATUS "Enabling C extension.") -define_extension_target( - _C - DESTINATION vllm - LANGUAGE ${VLLM_GPU_LANG} - SOURCES ${VLLM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} - ARCHITECTURES ${VLLM_GPU_ARCHES} - INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} - INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} - USE_SABI 3 - WITH_SOABI) - -# If CUTLASS is compiled on NVCC >= 12.5, it by default uses -# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the -# driver API. This causes problems when linking with earlier versions of CUDA. -# Setting this variable sidesteps the issue by calling the driver directly. -target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) + # If CUTLASS is compiled on NVCC >= 12.5, it by default uses + # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the + # driver API. This causes problems when linking with earlier versions of CUDA. + # Setting this variable sidesteps the issue by calling the driver directly. + target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) +endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # @@ -624,16 +387,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # set(VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/torch_bindings.cpp" + "csrc/libtorch_stable/cuda_view.cu" + "csrc/libtorch_stable/cuda_utils_kernels.cu" "csrc/libtorch_stable/activation_kernels.cu" + "csrc/libtorch_stable/ngram_embedding_kernels.cu" + "csrc/libtorch_stable/quantization/activation_kernels.cu" "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" "csrc/libtorch_stable/permute_cols.cu" "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" - "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu" "csrc/libtorch_stable/pos_encoding_kernels.cu" "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" + "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" "csrc/libtorch_stable/layernorm_kernels.cu" "csrc/libtorch_stable/layernorm_quant_kernels.cu" "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" @@ -642,18 +409,63 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/sampler.cu" "csrc/libtorch_stable/topk.cu" "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" - "csrc/libtorch_stable/attention/paged_attention_v1.cu" - "csrc/libtorch_stable/attention/paged_attention_v2.cu" - "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" "csrc/libtorch_stable/custom_all_reduce.cu" "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + if(VLLM_GPU_LANG STREQUAL "CUDA" AND + DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0) + + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS + "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1") + + endif() + endif() + if(VLLM_GPU_LANG STREQUAL "CUDA") + 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. + set(CUTLASS_REVISION "v4.4.2") + + # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided + if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR}) + set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR}) + endif() + + if(VLLM_CUTLASS_SRC_DIR) + if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR) + get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE) + endif() + message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation") + FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR}) + else() + FetchContent_Declare( + cutlass + GIT_REPOSITORY https://github.com/nvidia/cutlass.git + # Please keep this in sync with CUTLASS_REVISION line above. + GIT_TAG ${CUTLASS_REVISION} + GIT_PROGRESS TRUE + + # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. + # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. + # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE + GIT_SHALLOW TRUE + ) + endif() + FetchContent_MakeAvailable(cutlass) + list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" "csrc/libtorch_stable/cutlass_extensions/common.cpp" "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" @@ -661,22 +473,238 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") + # + # Machete kernels + # + # The machete kernels only work on hopper and require CUDA 12.0 or later. + # Only build Machete kernels if we are building for something compatible with sm90a + cuda_archs_loose_intersection(MACHETE_ARCHS "9.0a" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND MACHETE_ARCHS) + # + # For the Machete kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MACHETE_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py) + file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) + + message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") + message(STATUS "Last run machete generate script hash: $CACHE{MACHETE_GEN_SCRIPT_HASH}") + + if (NOT DEFINED CACHE{MACHETE_GEN_SCRIPT_HASH} + OR NOT $CACHE{MACHETE_GEN_SCRIPT_HASH} STREQUAL ${MACHETE_GEN_SCRIPT_HASH}) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MACHETE_GEN_SCRIPT} + RESULT_VARIABLE machete_generation_result + OUTPUT_VARIABLE machete_generation_output + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log + ) + + if (NOT machete_generation_result EQUAL 0) + message(FATAL_ERROR "Machete generation failed." + " Result: \"${machete_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/machete_generation.log") + else() + set(MACHETE_GEN_SCRIPT_HASH ${MACHETE_GEN_SCRIPT_HASH} + CACHE STRING "Last run machete generate script hash" FORCE) + message(STATUS "Machete generation completed successfully.") + endif() + else() + message(STATUS "Machete generation script has not changed, skipping generation.") + endif() + + # Add machete generated sources + file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu") + list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES}) + + # forward compatible + set_gencode_flags_for_srcs( + SRCS "${MACHETE_GEN_SOURCES}" + CUDA_ARCHS "${MACHETE_ARCHS}") + + list(APPEND VLLM_STABLE_EXT_SRC + csrc/libtorch_stable/quantization/machete/machete_pytorch.cu) + message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") + else() + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 + AND MACHETE_ARCHS) + message(STATUS "Not building Machete kernels as CUDA Compiler version is " + "not >= 12.0, we recommend upgrading to CUDA 12.0 or " + "later if you intend on running w4a16 quantized models on " + "Hopper.") + else() + message(STATUS "Not building Machete kernels as no compatible archs " + "found in CUDA target architectures") + endif() + endif() + set_gencode_flags_for_srcs( SRCS "${VLLM_STABLE_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") + if(COOPERATIVE_TOPK_ARCHS) + list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/cooperative_topk.cu") + set_gencode_flags_for_srcs( + SRCS "csrc/libtorch_stable/cooperative_topk.cu" + CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") + endif() + + # Only build Marlin kernels if we are building for at least some compatible archs. + # Keep building Marlin for 9.0 as there are some group sizes and shapes that + # are not supported by Machete yet. + + # marlin arches for fp16 output + # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0; + # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8). + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin has limited support for turing + cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}") + # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for fp8 input + # - sm80 doesn't support fp8 computation + # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction + # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for other files + cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") + + if (MARLIN_OTHER_ARCHS) + + # + # For the Marlin kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MARLIN_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) + list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") + + message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + + if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + RESULT_VARIABLE marlin_generation_result + OUTPUT_VARIABLE marlin_generation_result + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ) + + if (NOT marlin_generation_result EQUAL 0) + message(FATAL_ERROR "Marlin generation failed." + " Result: \"${marlin_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log") + else() + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + CACHE STRING "Last run Marlin generate script hash and arch" FORCE) + message(STATUS "Marlin generation completed successfully.") + endif() + else() + message(STATUS "Marlin generation script has not changed, skipping generation.") + endif() + + if (MARLIN_ARCHS) + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) + + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_BF16_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC}) + endif() + + if (MARLIN_SM75_ARCHS) + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_SM75_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC}) + endif() + + if (MARLIN_FP8_ARCHS) + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_FP8_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC}) + endif() + + set(MARLIN_SRCS + "csrc/libtorch_stable/quantization/marlin/marlin.cu" + "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" + "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_SRCS}" + CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_SRCS} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC "${MARLIN_SRCS}") + + message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}") + else() + message(STATUS "Not building Marlin kernels as no compatible archs found" + " in CUDA target architectures") + endif() + # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS) - set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") + set(DSV3_FUSED_A_GEMM_SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${DSV3_FUSED_A_GEMM_SRCS}" CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${DSV3_FUSED_A_GEMM_SRCS}") message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}") else() message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found " @@ -686,13 +714,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0. cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) - set(SRCS + set(FP32_ROUTER_GEMM_SRCS "csrc/libtorch_stable/fp32_router_gemm_entry.cu" "csrc/libtorch_stable/fp32_router_gemm.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${FP32_ROUTER_GEMM_SRCS}" CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${FP32_ROUTER_GEMM_SRCS}") message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}") else() message(STATUS "Not building fp32_router_gemm as no compatible archs found " @@ -702,13 +730,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Only build AllSpark kernels if we are building for at least some compatible archs. cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}") if (ALLSPARK_ARCHS) - set(SRCS + set(ALLSPARK_SRCS "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu" "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${ALLSPARK_SRCS}" CUDA_ARCHS "${ALLSPARK_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${ALLSPARK_SRCS}") message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}") else() message(STATUS "Not building AllSpark kernels as no compatible archs found" @@ -723,16 +751,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # CUDA 12.0 or later cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) - set(SRCS + set(SCALED_MM_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${SCALED_MM_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1") # Let scaled_mm_c2x know it doesn't need to build these arches list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") @@ -758,15 +786,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(SRCS + set(SCALED_MM_SM120_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" ) set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${SCALED_MM_SM120_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1") # Let scaled_mm_c2x know it doesn't need to build these arches list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") @@ -792,15 +820,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(SRCS + set(SCALED_MM_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" ) set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${SCALED_MM_SM100_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1") # Let scaled_mm_c2x know it doesn't need to build these arches list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}") @@ -826,11 +854,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # subtract out the archs that are already built for 3x list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS}) if (SCALED_MM_2X_ARCHS) - set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") + set(SCALED_MM_C2X_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${SCALED_MM_C2X_SRCS}" CUDA_ARCHS "${SCALED_MM_2X_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1") message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}") else() @@ -852,11 +880,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) - set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") + set(CUTLASS_MOE_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${CUTLASS_MOE_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1") message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}") else() @@ -876,11 +904,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") + set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${CUTLASS_MOE_SM100_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}") else() @@ -901,11 +929,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) - set(SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") + set(CUTLASS_MOE_DATA_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${CUTLASS_MOE_DATA_SRCS}" CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_DATA_SRCS}") message(STATUS "Building moe_data for archs: ${CUTLASS_MOE_DATA_ARCHS}") else() if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) @@ -922,71 +950,64 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch) # - # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require - # CUDA 12.8 or later + # SM12x FP4 kernels. These share some generic NVFP4 quantization entry + # sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends + # per-source flags, so shared files accumulate both SM12x and SM10x/11x + # gencodes when both families are requested. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") endif() - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS) - set(SRCS + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS) + set(FP4_SM120_SRCS "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" "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_sm120_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/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" - CUDA_ARCHS "${FP4_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") - set_gencode_flags_for_srcs( - SRCS "${NVFP4_KV_SRC}" - CUDA_ARCHS "${FP4_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") - target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) + SRCS "${FP4_SM120_SRCS}" + CUDA_ARCHS "${FP4_SM120_ARCHS}") + list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}") 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}") + message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}") else() - message(STATUS "Not building NVFP4 as no compatible archs were found.") - # clear FP4_ARCHS - set(FP4_ARCHS) + message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.") endif() - # FP4 Archs and flags + # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled + # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS) - set(SRCS + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS) + set(FP4_SM100_SRCS "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu") + "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + message(STATUS + "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") + endif() set_gencode_flags_for_srcs( - SRCS "${SRCS}" - CUDA_ARCHS "${FP4_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") - set_gencode_flags_for_srcs( - SRCS "${NVFP4_KV_SRC}" - CUDA_ARCHS "${FP4_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") - target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) + SRCS "${FP4_SM100_SRCS}" + CUDA_ARCHS "${FP4_SM100_ARCHS}") + list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}") 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}") + message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}") else() - message(STATUS "Not building NVFP4 as no compatible archs were found.") - # clear FP4_ARCHS - set(FP4_ARCHS) + message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.") endif() # @@ -996,17 +1017,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Only build W4A8 kernels if we are building for something compatible with sm90a cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS) - set(SRCS + set(W4A8_SRCS "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu" "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu" ) set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${W4A8_SRCS}" CUDA_ARCHS "${W4A8_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${W4A8_SRCS}") message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}") else() @@ -1022,22 +1043,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() endif() - # CUTLASS MLA Archs and flags + # CUTLASS MLA Archs and flags. + # Runtime dispatch is gated in + # vllm/v1/attention/backends/mla/cutlass_mla.py. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS) - set(SRCS + set(CUTLASS_MLA_SRCS "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${CUTLASS_MLA_SRCS}" CUDA_ARCHS "${MLA_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1") # Add MLA-specific include directories only to MLA source files - set_source_files_properties(${SRCS} + set_source_files_properties(${CUTLASS_MLA_SRCS} PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common") message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}") else() @@ -1049,11 +1072,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) - set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") + set(HADACORE_SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") set_gencode_flags_for_srcs( - SRCS "${SRCS}" + SRCS "${HADACORE_SRCS}" CUDA_ARCHS "${HADACORE_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${HADACORE_SRCS}") message(STATUS "Building hadacore") endif() @@ -1061,6 +1084,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() message(STATUS "Enabling C_stable extension.") + list(REMOVE_DUPLICATES VLLM_STABLE_EXT_SRC) define_extension_target( _C_stable_libtorch DESTINATION vllm @@ -1072,25 +1096,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") USE_SABI 3 WITH_SOABI) + # Set TORCH_TARGET_VERSION for stable ABI compatibility. + # This ensures we only use C-shim APIs available in PyTorch 2.11. + # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION + # which is currently set to 2.11. + target_compile_definitions(_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + # Needed to use cuda/hip APIs from C-shim if(VLLM_GPU_LANG STREQUAL "CUDA") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.11. - # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.11. - target_compile_definitions(_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020B000000000000ULL) target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA) + if(COOPERATIVE_TOPK_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_COOPERATIVE_TOPK=1) + endif() # Needed by CUTLASS kernels target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) elseif(VLLM_GPU_LANG STREQUAL "HIP") - # Set TORCH_TARGET_VERSION for stable ABI compatibility. - # This ensures we only use C-shim APIs available in PyTorch 2.10. - # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION - # which is currently set to 2.10. - target_compile_definitions(_C_stable_libtorch PRIVATE - TORCH_TARGET_VERSION=0x020A000000000000ULL) target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM) endif() @@ -1115,25 +1138,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # -# _moe_C extension +# _moe_C_stable_libtorch extension # set(VLLM_MOE_EXT_SRC - "csrc/moe/torch_bindings.cpp" - "csrc/moe/moe_align_sum_kernels.cu" - "csrc/moe/topk_softmax_kernels.cu" - "csrc/moe/topk_softplus_sqrt_kernels.cu") + "csrc/libtorch_stable/moe/torch_bindings.cpp" + "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu" + "csrc/libtorch_stable/moe/topk_softmax_kernels.cu" + "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC - "csrc/moe/moe_wna16.cu" - "csrc/moe/grouped_topk_kernels.cu") + "csrc/libtorch_stable/moe/moe_wna16.cu" + "csrc/libtorch_stable/moe/grouped_topk_kernels.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") set(MOE_PERMUTE_SRC - "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" - "csrc/moe/moe_permute_unpermute_op.cu") + "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" + "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu") list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}") endif() @@ -1144,7 +1167,7 @@ set_gencode_flags_for_srcs( if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_MOE_WNA16_SRC - "csrc/moe/moe_wna16.cu") + "csrc/libtorch_stable/moe/moe_wna16.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_MOE_WNA16_SRC}" @@ -1165,7 +1188,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # - sm80 doesn't support fp8 computation # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) - cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") + endif() # moe marlin arches for other files cuda_archs_loose_intersection(MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") if (MARLIN_MOE_OTHER_ARCHS) @@ -1175,7 +1202,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # preselected input type pairs and schedules. # Generate sources: set(MOE_MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py) file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -1210,7 +1237,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_ARCHS) - file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu") + file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SRC}" CUDA_ARCHS "${MARLIN_MOE_ARCHS}") @@ -1222,7 +1249,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_SM75_ARCHS) - file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu") + file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SM75_SRC}" CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}") @@ -1234,7 +1261,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_FP8_ARCHS) - file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu") + file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_FP8_SRC}" CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}") @@ -1245,7 +1272,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC}) endif() - set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu") + set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_OTHER_SRC}" CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}") @@ -1266,9 +1293,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS) set(DSV3_ROUTER_GEMM_SRC - "csrc/moe/dsv3_router_gemm_entry.cu" - "csrc/moe/dsv3_router_gemm_float_out.cu" - "csrc/moe/dsv3_router_gemm_bf16_out.cu") + "csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu" + "csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu" + "csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu") set_gencode_flags_for_srcs( SRCS "${DSV3_ROUTER_GEMM_SRC}" CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}") @@ -1281,9 +1308,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() endif() -message(STATUS "Enabling moe extension.") +message(STATUS "Enabling MoE C_stable extension.") define_extension_target( - _moe_C + _moe_C_stable_libtorch DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_MOE_EXT_SRC} @@ -1294,6 +1321,42 @@ define_extension_target( USE_SABI 3 WITH_SOABI) +# Set TORCH_TARGET_VERSION for stable ABI compatibility. +# This ensures we only use C-shim APIs available in PyTorch 2.11. +# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION +# which is currently set to 2.11. +target_compile_definitions(_moe_C_stable_libtorch PRIVATE + TORCH_TARGET_VERSION=0x020B000000000000ULL) + +# Needed to use cuda/hip APIs from C-shim +if(VLLM_GPU_LANG STREQUAL "CUDA") + target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA) + # Needed by CUTLASS kernels + target_compile_definitions(_moe_C_stable_libtorch PRIVATE + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) +elseif(VLLM_GPU_LANG STREQUAL "HIP") + target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM) +endif() + +# On ROCm, _moe_C_stable_libtorch calls raw HIP APIs (e.g. hipGetDevice in +# get_device_prop()) which must resolve to the same libamdhip64.so that +# PyTorch uses. When PyTorch bundles its own copy (pip/conda wheels), +# the raw HIP calls would otherwise resolve to the system ROCm copy, +# initializing a second HIP runtime that corrupts device state (wrong +# device on DeviceGuard, core dumps on multi-GPU tests). +# +# If PyTorch doesn't bundle libamdhip64 (built from source against system +# ROCm), there is only one copy in the process and no action is needed — +# the HIP compiler already links the system libamdhip64 automatically. +if(VLLM_GPU_LANG STREQUAL "HIP") + find_library(_MOE_STABLE_TORCH_AMDHIP64 amdhip64 + PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH) + if(_MOE_STABLE_TORCH_AMDHIP64) + message(STATUS "Found PyTorch-bundled libamdhip64 for _moe_C_stable_libtorch at ${_MOE_STABLE_TORCH_AMDHIP64}") + target_link_libraries(_moe_C_stable_libtorch PRIVATE ${_MOE_STABLE_TORCH_AMDHIP64}) + endif() +endif() + if(VLLM_GPU_LANG STREQUAL "HIP") # # _rocm_C extension @@ -1341,6 +1404,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/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) include(cmake/external_projects/qutlass.cmake) diff --git a/MANIFEST.in b/MANIFEST.in index fb3cccbb4a9..cbb136e6b76 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -4,6 +4,7 @@ include requirements/cuda.txt include requirements/rocm.txt include requirements/cpu.txt include CMakeLists.txt +include tools/build_rust.py recursive-include cmake * recursive-include csrc * diff --git a/SECURITY.md b/SECURITY.md index d6319cdb1ac..1e2a5a0adef 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -34,6 +34,15 @@ Vulnerabilities that cause denial of service or partial disruption, but do not a Minor issues such as informational disclosures, logging errors, non-exploitable flaws, or weaknesses that require local or high-privilege access and offer negligible impact. Examples include side channel attacks or hash collisions. These issues often have CVSS scores less than 4.0 +## Fix disclosure policy + +When a security report is accepted, the fix process depends on the severity: + +* **CRITICAL and HIGH severity**: Fixes are developed in a private security fork and coordinated with the prenotification group before public disclosure. +* **MODERATE and LOW severity**: Fixes are developed and submitted as public pull requests. These issues do not require embargo since they do not enable arbitrary code execution or significant data breach, and public visibility accelerates community review and adoption of the fix. + +The vulnerability management team reserves the right to adjust the disclosure approach on a case-by-case basis, taking into account factors such as active exploitation, unusual attack surface, or coordination requirements with downstream vendors. + ## Prenotification policy For certain security issues of CRITICAL, HIGH, or MODERATE severity level, we may prenotify certain organizations or vendors that ship vLLM. The purpose of this prenotification is to allow for a coordinated release of fixes for severe issues. diff --git a/benchmarks/attention_benchmarks/README.md b/benchmarks/attention_benchmarks/README.md index afce3443316..944ceb91af9 100644 --- a/benchmarks/attention_benchmarks/README.md +++ b/benchmarks/attention_benchmarks/README.md @@ -108,7 +108,6 @@ python benchmark.py \ --backends flash triton flashinfer \ --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \ --num-layers 10 \ - --repeats 5 \ --output-csv results.csv ``` @@ -164,14 +163,17 @@ python benchmark.py \ # Model configuration --num-layers N # Number of layers --head-dim N # Head dimension +--v-head-dim N # Value head dimension (defaults to --head-dim) --num-q-heads N # Query heads --num-kv-heads N # KV heads --block-size N # Block size +--kv-lora-rank N # MLA KV LoRA rank +--qk-nope-head-dim N # MLA non-RoPE QK head dim +--qk-rope-head-dim N # MLA RoPE QK head dim # Benchmark settings --device DEVICE # Device (default: cuda:0) ---repeats N # Repetitions ---warmup-iters N # Warmup iterations +--warmup-ms N # Warmup window in ms for triton do_bench --profile-memory # Profile memory usage # Parameter sweeps @@ -211,8 +213,6 @@ config = BenchmarkConfig( num_kv_heads=1, block_size=128, device="cuda:0", - repeats=5, - warmup_iters=3, ) # CUTLASS MLA with specific num_kv_splits @@ -253,14 +253,10 @@ formatter.save_json(results, "output.json") ## Tips -**1. Warmup matters** - Use `--warmup-iters 10` for stable results +**1. Save results** - Always use `--output-csv` or `--output-json` -**2. Multiple repeats** - Use `--repeats 20` for low variance +**2. Test incrementally** - Start with `--num-layers 1` -**3. Save results** - Always use `--output-csv` or `--output-json` +**3. Extended grammar** - Leverage spec decode, chunked prefill patterns -**4. Test incrementally** - Start with `--num-layers 1 --repeats 1` - -**5. Extended grammar** - Leverage spec decode, chunked prefill patterns - -**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values +**4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index c4c331f7f8e..9860d4b2d1c 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -26,6 +26,9 @@ Examples: """ import argparse +import os +import shutil +import subprocess import sys from dataclasses import replace from pathlib import Path @@ -50,6 +53,16 @@ from common import ( from vllm.v1.worker.workspace import init_workspace_manager +def _str2bool(v) -> bool: + if isinstance(v, bool): + return v + if v.lower() in ("true", "1", "yes", "t"): + return True + if v.lower() in ("false", "0", "no", "f"): + return False + raise argparse.ArgumentTypeError(f"expected a boolean, got {v!r}") + + def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: """Run standard attention benchmark (Flash/Triton/FlashInfer).""" from runner import run_attention_benchmark @@ -83,13 +96,15 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult: else: return run_standard_attention_benchmark(config) except Exception as e: + error_msg = str(e) or repr(e) return BenchmarkResult( config=config, mean_time=float("inf"), + median_time=float("inf"), std_time=0, min_time=float("inf"), max_time=float("inf"), - error=str(e), + error=error_msg, ) @@ -115,9 +130,12 @@ def run_model_parameter_sweep( """ all_results = [] - console.print( - f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]" + sweep_desc = ( + f"{sweep.param_name} = {sweep.values}" + if sweep.param_name + else f"{len(sweep.values)} configurations" ) + console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]") total = len(backends) * len(batch_specs) * len(sweep.values) @@ -125,9 +143,9 @@ def run_model_parameter_sweep( for backend in backends: for spec in batch_specs: for value in sweep.values: - # Create config with modified model parameter + # Create config with modified model parameter(s) config_args = base_config_args.copy() - config_args[sweep.param_name] = value + sweep.apply(config_args, value) # Create config with original backend for running clean_config = BenchmarkConfig( @@ -144,13 +162,21 @@ def run_model_parameter_sweep( all_results.append(result) if not result.success: + err_label = ( + f"{sweep.param_name}={value}" + if sweep.param_name + else f"{value}" + ) console.print( - f"[red]Error {backend} {spec} {sweep.param_name}=" - f"{value}: {result.error}[/]" + f"[red]Error {backend} {spec} {err_label}" + f": {result.error}[/]" ) pbar.update(1) + if base_config_args.get("ncu_profile"): + return all_results + # Display sweep results - create separate table for each parameter value console.print("\n[bold green]Model Parameter Sweep Results:[/]") formatter = ResultsFormatter(console) @@ -184,7 +210,10 @@ def run_model_parameter_sweep( ) for param_value in sorted_param_values: - console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]") + label = ( + f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value + ) + console.print(f"\n[bold cyan]{label}[/]") param_results = by_param_value[param_value] # Create modified results with original backend names @@ -200,8 +229,9 @@ def run_model_parameter_sweep( formatter.print_table(modified_results, backends, compare_to_fastest=True) # Show optimal backend for each (param_value, batch_spec) combination + sweep_name = sweep.param_name or "config" console.print( - f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]" + f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]" ) # Group by (param_value, batch_spec) @@ -236,7 +266,10 @@ def run_model_parameter_sweep( for param_value, spec in sorted_keys: # Print header when param value changes if param_value != current_param_value: - console.print(f"\n [bold]{sweep.param_name}={param_value}:[/]") + header = ( + f"{sweep.param_name}={param_value}" if sweep.param_name else param_value + ) + console.print(f"\n [bold]{header}:[/]") current_param_value = param_value results = by_param_and_spec[(param_value, spec)] @@ -322,6 +355,9 @@ def run_parameter_sweep( pbar.update(1) + if base_config_args.get("ncu_profile"): + return all_results + # Display sweep results console.print("\n[bold green]Sweep Results:[/]") backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values] @@ -459,6 +495,20 @@ def main(): help="Prefill backends to compare (fa2, fa3, fa4). " "Uses the first decode backend for impl construction.", ) + parser.add_argument( + "--fp8-output-scale", + type=float, + help="Static per-tensor scale enabling the MLA prefill FP8-output " + "comparison on FA4 (fused write vs standalone post-quant).", + ) + parser.add_argument( + "--fuse-quant-op", + nargs="+", + type=_str2bool, + help="FP8-output write path(s) to run: false = bf16 attention + " + "standalone static-FP8 quant, true = FA4 writes FP8 directly. " + "Default: both.", + ) # Batch specifications parser.add_argument( @@ -474,11 +524,35 @@ def main(): parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads") parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads") parser.add_argument("--block-size", type=int, default=16, help="Block size") + parser.add_argument( + "--v-head-dim", + type=int, + default=None, + help="Value head dimension (defaults to --head-dim if unset)", + ) + + # MLA-specific model dimensions + parser.add_argument( + "--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank" + ) + parser.add_argument( + "--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim" + ) + parser.add_argument( + "--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim" + ) # Benchmark settings parser.add_argument("--device", default="cuda:0", help="Device") - parser.add_argument("--repeats", type=int, default=1, help="Repetitions") - parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations") + parser.add_argument( + "--warmup-ms", + type=int, + default=None, + help=( + "Warmup window in ms for triton's do_bench (default: triton's own). " + "Has no effect with CUDA graphs; pass --no-cuda-graphs to use it." + ), + ) parser.add_argument("--profile-memory", action="store_true", help="Profile memory") parser.add_argument( "--kv-cache-dtype", @@ -491,10 +565,33 @@ def main(): action=argparse.BooleanOptionalAction, default=True, help=( - "Launch kernels with CUDA graphs to eliminate CPU overhead" - "in measurements (default: True)" + "Use triton do_bench_cudagraph (True) or do_bench (False) " + "for timing. CUDA graphs eliminate CPU launch overhead " + "(default: True)" ), ) + parser.add_argument( + "--num-splits", + type=int, + default=None, + help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)", + ) + parser.add_argument( + "--ncu-profile", + action="store_true", + default=False, + help=( + "Enable Nsight Compute profiling mode. Automatically wraps the " + "script with ncu, capturing a profile with source correlation. " + "Use --ncu-output to set the output file name." + ), + ) + parser.add_argument( + "--ncu-output", + type=str, + default="profile", + help="Output file name for ncu profile (default: 'profile').", + ) # Parameter sweep (use YAML config for advanced sweeps) parser.add_argument( @@ -545,6 +642,12 @@ def main(): # Prefill backends (e.g., ["fa3", "fa4"]) args.prefill_backends = yaml_config.get("prefill_backends", None) + # FP8 output benchmark knobs; CLI wins. + if args.fp8_output_scale is None: + args.fp8_output_scale = yaml_config.get("fp8_output_scale", None) + if args.fuse_quant_op is None: + args.fuse_quant_op = yaml_config.get("fuse_quant_op", None) + # Check for special modes args.mode = yaml_config.get("mode", None) @@ -576,23 +679,28 @@ def main(): model = yaml_config["model"] args.num_layers = model.get("num_layers", args.num_layers) args.head_dim = model.get("head_dim", args.head_dim) + args.v_head_dim = model.get("v_head_dim", args.v_head_dim) args.num_q_heads = model.get("num_q_heads", args.num_q_heads) args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads) args.block_size = model.get("block_size", args.block_size) + # MLA-specific dimensions + args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank) + args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim) + args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim) # Benchmark settings (top-level keys) if "device" in yaml_config: args.device = yaml_config["device"] - if "repeats" in yaml_config: - args.repeats = yaml_config["repeats"] - if "warmup_iters" in yaml_config: - args.warmup_iters = yaml_config["warmup_iters"] + if "warmup_ms" in yaml_config: + args.warmup_ms = yaml_config["warmup_ms"] if "profile_memory" in yaml_config: args.profile_memory = yaml_config["profile_memory"] if "kv_cache_dtype" in yaml_config: args.kv_cache_dtype = yaml_config["kv_cache_dtype"] if "cuda_graphs" in yaml_config: args.cuda_graphs = yaml_config["cuda_graphs"] + if "ncu_profile" in yaml_config: + args.ncu_profile = yaml_config["ncu_profile"] # Parameter sweep configuration if "parameter_sweep" in yaml_config: @@ -612,7 +720,7 @@ def main(): if "model_parameter_sweep" in yaml_config: sweep_config = yaml_config["model_parameter_sweep"] args.model_parameter_sweep = ModelParameterSweep( - param_name=sweep_config["param_name"], + param_name=sweep_config.get("param_name"), values=sweep_config["values"], label_format=sweep_config.get( "label_format", "{backend}_{param_name}_{value}" @@ -631,6 +739,32 @@ def main(): console.print() + # Re-exec under ncu if --ncu-profile and not already inside ncu. This runs + # after YAML processing so ncu_profile set via config file is honored. + if args.ncu_profile and "_NCU_INNER" not in os.environ: + ncu = shutil.which("ncu") + if ncu is None: + print("Error: 'ncu' not found in PATH", file=sys.stderr) + sys.exit(1) + cmd = [ + ncu, + "--profile-from-start", + "off", + "--set", + "full", + "--import-source", + "yes", + "-o", + args.ncu_output, + sys.executable, + *sys.argv, + ] + env = os.environ.copy() + env["CUTE_DSL_LINEINFO"] = "1" + env["_NCU_INNER"] = "1" + print(f"Launching: {' '.join(cmd)}") + sys.exit(subprocess.call(cmd, env=env)) + # Handle CLI-based parameter sweep (if not from YAML) if ( (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None) @@ -655,6 +789,18 @@ def main(): console.print(f"Batch specs: {', '.join(args.batch_specs)}") console.print(f"KV cache dtype: {args.kv_cache_dtype}") console.print(f"CUDA graphs: {args.cuda_graphs}") + if args.warmup_ms is not None and args.cuda_graphs: + console.print( + "[yellow]Warning: --warmup-ms is ignored with CUDA graphs " + "(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs " + "to use it.[/]" + ) + if args.num_splits == 0 and args.cuda_graphs: + console.print( + "[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph " + "compatible and may fail or fall back. Pass --no-cuda-graphs or use " + "--num-splits >=1.[/]" + ) console.print() init_workspace_manager(args.device) @@ -662,8 +808,68 @@ def main(): # Run benchmarks all_results = [] + # Under ncu profiling the kernels run only to be captured by the profiler; + # timings are placeholder zeros, so the result tables and saved metrics are + # skipped. The Nsight Compute report (--ncu-output) holds the real data. + if args.ncu_profile: + console.print( + "[dim]ncu profiling enabled: result tables and saved metrics are " + "skipped (timings are placeholder zeros).[/]" + ) + + # FA4 fused FP8 output vs standalone post-quant, on the same fa4 kernel: + # the delta is the post-quant kernel the fused path removes. + fp8_output_scale = getattr(args, "fp8_output_scale", None) + if fp8_output_scale is not None: + decode_backend = backends[0] + fuse_variants = args.fuse_quant_op or [False, True] + label_of = {False: "post_quant", True: "fused"} + console.print( + f"[yellow]FP8 output comparison @ scale={fp8_output_scale} " + f"(prefill=fa4, decode impl={decode_backend})[/]" + ) + fp8_results = [] + total = len(fuse_variants) * len(args.batch_specs) + with tqdm(total=total, desc="FP8 output benchmarking") as pbar: + for spec in args.batch_specs: + for fuse in fuse_variants: + config = BenchmarkConfig( + backend=decode_backend, + batch_spec=spec, + num_layers=args.num_layers, + head_dim=args.head_dim, + num_q_heads=args.num_q_heads, + num_kv_heads=args.num_kv_heads, + block_size=args.block_size, + device=args.device, + repeats=args.repeats, + warmup_iters=args.warmup_iters, + profile_memory=args.profile_memory, + kv_cache_dtype=args.kv_cache_dtype, + use_cuda_graphs=args.cuda_graphs, + prefill_backend="fa4", + ) + result = run_benchmark( + config, output_scale=fp8_output_scale, fuse_quant_op=fuse + ) + label = label_of[fuse] + labeled_config = replace(result.config, backend=label) + result = replace(result, config=labeled_config) + fp8_results.append(result) + + if not result.success: + console.print(f"[red]Error {label} {spec}: {result.error}[/]") + + pbar.update(1) + + console.print("\n[bold green]FP8 Output Results:[/]") + formatter = ResultsFormatter(console) + labels = [label_of[f] for f in fuse_variants] + formatter.print_table(fp8_results, labels, compare_to_fastest=True) + all_results = fp8_results + # Handle special mode: decode_vs_prefill comparison - if hasattr(args, "mode") and args.mode == "decode_vs_prefill": + elif hasattr(args, "mode") and args.mode == "decode_vs_prefill": console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]") console.print( "[dim]For each query length, testing both decode and prefill pipelines[/]" @@ -708,11 +914,11 @@ def main(): num_kv_heads=args.num_kv_heads, block_size=args.block_size, device=args.device, - repeats=args.repeats, - warmup_iters=args.warmup_iters, profile_memory=args.profile_memory, kv_cache_dtype=args.kv_cache_dtype, use_cuda_graphs=args.cuda_graphs, + ncu_profile=args.ncu_profile, + warmup_ms=args.warmup_ms, ) # Add decode pipeline config @@ -749,6 +955,7 @@ def main(): result = BenchmarkResult( config=config, mean_time=timing["mean"], + median_time=timing.get("median", timing["mean"]), std_time=timing["std"], min_time=timing["min"], max_time=timing["max"], @@ -770,6 +977,7 @@ def main(): result = BenchmarkResult( config=config, mean_time=float("inf"), + median_time=float("inf"), std_time=0, min_time=float("inf"), max_time=float("inf"), @@ -779,6 +987,9 @@ def main(): pbar.update(1) + if args.ncu_profile: + return + # Display decode vs prefill results console.print("\n[bold green]Decode vs Prefill Results:[/]") @@ -858,15 +1069,20 @@ def main(): base_config_args = { "num_layers": args.num_layers, "head_dim": args.head_dim, + "v_head_dim": args.v_head_dim, "num_q_heads": args.num_q_heads, "num_kv_heads": args.num_kv_heads, "block_size": args.block_size, "device": args.device, - "repeats": args.repeats, - "warmup_iters": args.warmup_iters, "profile_memory": args.profile_memory, "kv_cache_dtype": args.kv_cache_dtype, "use_cuda_graphs": args.cuda_graphs, + "ncu_profile": args.ncu_profile, + "warmup_ms": args.warmup_ms, + "num_splits": args.num_splits, + "kv_lora_rank": args.kv_lora_rank, + "qk_nope_head_dim": args.qk_nope_head_dim, + "qk_rope_head_dim": args.qk_rope_head_dim, } all_results = run_model_parameter_sweep( backends, @@ -882,15 +1098,17 @@ def main(): base_config_args = { "num_layers": args.num_layers, "head_dim": args.head_dim, + "v_head_dim": args.v_head_dim, "num_q_heads": args.num_q_heads, "num_kv_heads": args.num_kv_heads, "block_size": args.block_size, "device": args.device, - "repeats": args.repeats, - "warmup_iters": args.warmup_iters, "profile_memory": args.profile_memory, "kv_cache_dtype": args.kv_cache_dtype, "use_cuda_graphs": args.cuda_graphs, + "ncu_profile": args.ncu_profile, + "warmup_ms": args.warmup_ms, + "num_splits": args.num_splits, } all_results = run_parameter_sweep( backends, args.batch_specs, base_config_args, args.parameter_sweep, console @@ -914,15 +1132,17 @@ def main(): batch_spec=spec, num_layers=args.num_layers, head_dim=args.head_dim, + v_head_dim=getattr(args, "v_head_dim", None), num_q_heads=args.num_q_heads, num_kv_heads=args.num_kv_heads, block_size=args.block_size, device=args.device, - repeats=args.repeats, - warmup_iters=args.warmup_iters, profile_memory=args.profile_memory, kv_cache_dtype=args.kv_cache_dtype, use_cuda_graphs=args.cuda_graphs, + ncu_profile=args.ncu_profile, + warmup_ms=args.warmup_ms, + num_splits=args.num_splits, ) result = run_benchmark(config) @@ -935,9 +1155,10 @@ def main(): pbar.update(1) - console.print("\n[bold green]Results:[/]") - formatter = ResultsFormatter(console) - formatter.print_table(decode_results, backends) + if not args.ncu_profile: + console.print("\n[bold green]Results:[/]") + formatter = ResultsFormatter(console) + formatter.print_table(decode_results, backends) # Run prefill backend comparison if prefill_backends: @@ -962,9 +1183,8 @@ def main(): num_kv_heads=args.num_kv_heads, block_size=args.block_size, device=args.device, - repeats=args.repeats, - warmup_iters=args.warmup_iters, profile_memory=args.profile_memory, + warmup_ms=args.warmup_ms, prefill_backend=pb, ) @@ -980,16 +1200,17 @@ def main(): pbar.update(1) - console.print("\n[bold green]Prefill Backend Results:[/]") - formatter = ResultsFormatter(console) - formatter.print_table( - prefill_results, prefill_backends, compare_to_fastest=True - ) + if not args.ncu_profile: + console.print("\n[bold green]Prefill Backend Results:[/]") + formatter = ResultsFormatter(console) + formatter.print_table( + prefill_results, prefill_backends, compare_to_fastest=True + ) all_results = decode_results + prefill_results - # Save results - if all_results: + # Save results (skip ncu profiling runs: timings are placeholder zeros) + if all_results and not args.ncu_profile: formatter = ResultsFormatter(console) if args.output_csv: formatter.save_csv(all_results, args.output_csv) diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 74d9e239725..106d7854804 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -15,6 +15,8 @@ from batch_spec import get_batch_type, parse_batch_spec from rich.console import Console from rich.table import Table +from vllm.triton_utils import triton + def batch_spec_sort_key(spec: str) -> tuple[int, int, int]: """ @@ -34,6 +36,30 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]: return (0, 0, 0) +def run_do_bench( + benchmark_fn, + use_cuda_graphs: bool, + warmup_ms: int | None = None, +) -> list[float]: + kwargs: dict[str, Any] = {"return_mode": "all"} + if use_cuda_graphs: + result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs) + else: + if warmup_ms is not None: + kwargs["warmup"] = warmup_ms + result = triton.testing.do_bench(benchmark_fn, **kwargs) + return result + + +def run_ncu_profile(benchmark_fn) -> None: + benchmark_fn() + torch.accelerator.synchronize() + torch.cuda.cudart().cudaProfilerStart() + benchmark_fn() + torch.accelerator.synchronize() + torch.cuda.cudart().cudaProfilerStop() + + # Mock classes for vLLM attention infrastructure @@ -182,18 +208,37 @@ class ParameterSweep: @dataclass class ModelParameterSweep: - """Configuration for sweeping a model configuration parameter.""" + """Configuration for sweeping model configuration parameter(s). - param_name: str # Name of the model config parameter to sweep (e.g., "num_q_heads") - values: list[Any] # List of values to test - label_format: str = "{backend}_{param_name}_{value}" # Result label template + Supports two modes: + - Single param: param_name="head_dim", values=[128, 256, 512] + - Multi param: values=[{head_dim: 192, v_head_dim: 128}, {head_dim: 256}] + When values are dicts, each dict's keys are applied as config overrides. + """ + + param_name: str | None = None + values: list[Any] | None = None + label_format: str = "{backend}_{param_name}_{value}" def get_label(self, backend: str, value: Any) -> str: """Generate a label for a specific parameter value.""" + if isinstance(value, dict): + return self.label_format.format( + backend=backend, param_name=self.param_name, value=value, **value + ) return self.label_format.format( backend=backend, param_name=self.param_name, value=value ) + def apply(self, config_args: dict, value: Any) -> None: + """Apply a sweep value to config args.""" + if isinstance(value, dict): + config_args.update(value) + elif self.param_name is not None: + config_args[self.param_name] = value + else: + raise ValueError("param_name must be set if sweep values are not dicts") + @dataclass class BenchmarkConfig: @@ -208,10 +253,10 @@ class BenchmarkConfig: block_size: int device: str dtype: torch.dtype = torch.float16 - repeats: int = 1 - warmup_iters: int = 3 profile_memory: bool = False use_cuda_graphs: bool = False + ncu_profile: bool = False + warmup_ms: int | None = None # "auto" or "fp8" kv_cache_dtype: str = "auto" @@ -226,6 +271,7 @@ class BenchmarkConfig: # Backend-specific tuning num_kv_splits: int | None = None # CUTLASS MLA reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA + num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled) @dataclass @@ -234,6 +280,7 @@ class BenchmarkResult: config: BenchmarkConfig mean_time: float # seconds + median_time: float # seconds std_time: float # seconds min_time: float # seconds max_time: float # seconds @@ -252,6 +299,7 @@ class BenchmarkResult: return { "config": asdict(self.config), "mean_time": self.mean_time, + "median_time": self.median_time, "std_time": self.std_time, "min_time": self.min_time, "max_time": self.max_time, diff --git a/benchmarks/attention_benchmarks/configs/mla_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_decode.yaml index 8f12ac72306..c1d47bf5748 100644 --- a/benchmarks/attention_benchmarks/configs/mla_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_decode.yaml @@ -56,8 +56,6 @@ backends: - TOKENSPEED_MLA # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8) device: "cuda:0" -repeats: 100 -warmup_iters: 10 profile_memory: true # Backend-specific tuning diff --git a/benchmarks/attention_benchmarks/configs/mla_fa4_fp8_output.yaml b/benchmarks/attention_benchmarks/configs/mla_fa4_fp8_output.yaml new file mode 100644 index 00000000000..85588fcf958 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/mla_fa4_fp8_output.yaml @@ -0,0 +1,44 @@ +# MLA prefill FP8-output microbenchmark (FA4). +# Compares the fused FP8 write against bf16 attention + a standalone static-FP8 +# quant; the delta is the post-quant kernel the fused path removes. +# DeepSeek-Coder-V2-Lite dims; FA4 needs SM100/110. +# +# Usage: +# python benchmark.py --config configs/mla_fa4_fp8_output.yaml + +description: "MLA prefill FA4 fused-FP8 output vs post-quant" + +model: + name: "deepseek-v2-lite" + num_layers: 27 + num_q_heads: 16 + num_kv_heads: 1 + head_dim: 576 + kv_lora_rank: 512 + qk_nope_head_dim: 128 + qk_rope_head_dim: 64 + v_head_dim: 128 + block_size: 128 + +# Pure prefill (q_len == kv_len) so every token goes through forward_mha. +batch_specs: + - "q512" + - "q1k" + - "q2k" + - "q4k" + - "q8k" + - "2q4k" + - "4q4k" + - "8q4k" + +# Only used to construct the MLA impl; the pure-prefill specs skip decode. +decode_backends: + - CUTLASS_MLA + +# Sweep the two FP8 write paths (prefill backend is fixed to fa4). +fp8_output_scale: 0.1 +fuse_quant_op: [false, true] + +device: "cuda:0" +repeats: 50 +warmup_iters: 10 diff --git a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml index c342e9fb8c1..fcb1d8639b7 100644 --- a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml @@ -51,8 +51,6 @@ backends: - FLASHMLA # Hopper only device: "cuda:0" -repeats: 5 -warmup_iters: 3 profile_memory: true # Analyze chunked prefill workspace size impact diff --git a/benchmarks/attention_benchmarks/configs/mla_prefill.yaml b/benchmarks/attention_benchmarks/configs/mla_prefill.yaml index 1e1ab264bac..f39cdd8d1c2 100644 --- a/benchmarks/attention_benchmarks/configs/mla_prefill.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_prefill.yaml @@ -124,5 +124,3 @@ prefill_backends: - tokenspeed device: "cuda:0" -repeats: 20 -warmup_iters: 5 diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml index 689c9f3c3c6..c791638241f 100644 --- a/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml @@ -53,6 +53,4 @@ backends: - FLASHINFER_MLA_SPARSE device: "cuda:0" -repeats: 100 -warmup_iters: 10 profile_memory: true diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml index ef6b2cb07dc..fd8a0e22c5e 100644 --- a/benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml @@ -57,6 +57,4 @@ backends: - FLASHINFER_MLA_SPARSE device: "cuda:0" -repeats: 10 -warmup_iters: 3 profile_memory: true diff --git a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml index 0d76ef0a358..9f53eac2c9c 100644 --- a/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml +++ b/benchmarks/attention_benchmarks/configs/reorder_threshold.yaml @@ -63,8 +63,6 @@ model: # Benchmark settings device: "cuda:0" -repeats: 15 # More repeats for spec decode variance -warmup_iters: 5 profile_memory: false # Output diff --git a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml index 47b6d3604d1..5e8775f0a42 100644 --- a/benchmarks/attention_benchmarks/configs/speculative_decode.yaml +++ b/benchmarks/attention_benchmarks/configs/speculative_decode.yaml @@ -49,8 +49,6 @@ backends: # Benchmark settings device: "cuda:0" -repeats: 10 # More repeats for statistical significance -warmup_iters: 5 profile_memory: false # Test these threshold values for optimization diff --git a/benchmarks/attention_benchmarks/configs/standard_attention.yaml b/benchmarks/attention_benchmarks/configs/standard_attention.yaml index deb5a4b27ff..ccd44a426b9 100644 --- a/benchmarks/attention_benchmarks/configs/standard_attention.yaml +++ b/benchmarks/attention_benchmarks/configs/standard_attention.yaml @@ -43,6 +43,4 @@ backends: - FLASHINFER device: "cuda:0" -repeats: 5 -warmup_iters: 3 profile_memory: false diff --git a/benchmarks/attention_benchmarks/configs/standard_decode.yaml b/benchmarks/attention_benchmarks/configs/standard_decode.yaml new file mode 100644 index 00000000000..0861bd63dad --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/standard_decode.yaml @@ -0,0 +1,142 @@ +# Standard attention decode benchmark configuration +# Sweeps num_q_heads and num_kv_heads to isolate effects of: +# 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads) +# 2. Absolute head count (fixed 4:1 ratio, vary scale) + +model: + num_layers: 32 + num_q_heads: 32 # Base value, overridden by sweep + num_kv_heads: 8 # Base value, overridden by sweep + head_dim: 128 + block_size: 16 + +# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and +# head_dim where it differs from the base (128). Head counts are per-GPU +# (i.e. after TP sharding). +# +# Group A — vary GQA ratio (fixed q=32, head_dim=128): +# 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA) +# +# Groups B-E — real model configs at various TP degrees: +# Model head_dim Full TP2 TP4 TP8 +# Llama 3 8B 128 32:8 16:4 8:2 4:1 +# Llama 3 70B 128 64:8 32:4 16:2 8:1 +# GPT-OSS 120B 64 64:8 32:4 16:2 8:1 +# Llama 3 405B 128 128:8 64:4 32:2 16:1 +model_parameter_sweep: + values: + # --- head_dim=128 (Llama 3 family) --- + - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1 + - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1 + - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8 + - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4 + - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2 + - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1 + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1 + - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8 + - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4 + - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2 + - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1 + # --- head_dim=64 (GPT-OSS 120B) --- + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1 + label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}" + +batch_specs: + # ---- batch_size x seq_len grid (decode: q_len=1) ---- + # Small grid for quick iteration. Uncomment for full sweep. + + # Batch size 1 + - "q1s1k" + - "q1s512" + - "q1s2k" + - "q1s4k" + - "q1s8k" + - "q1s16k" + - "q1s32k" + + # Batch size 2 + - "2q1s512" + - "2q1s1k" + - "2q1s2k" + - "2q1s4k" + - "2q1s8k" + - "2q1s16k" + - "2q1s32k" + + # Batch size 4 + - "4q1s512" + - "4q1s1k" + - "4q1s2k" + - "4q1s4k" + - "4q1s8k" + - "4q1s16k" + - "4q1s32k" + + # Batch size 8 + - "8q1s1k" + - "8q1s512" + - "8q1s2k" + - "8q1s4k" + - "8q1s8k" + - "8q1s16k" + - "8q1s32k" + + # Batch size 16 + - "16q1s512" + - "16q1s1k" + - "16q1s2k" + - "16q1s4k" + - "16q1s8k" + - "16q1s16k" + - "16q1s32k" + + # Batch size 32 + - "32q1s512" + - "32q1s1k" + - "32q1s2k" + - "32q1s4k" + - "32q1s8k" + - "32q1s16k" + - "32q1s32k" + + # Batch size 64 + - "64q1s1k" + - "64q1s512" + - "64q1s2k" + - "64q1s4k" + - "64q1s8k" + - "64q1s16k" + - "64q1s32k" + + # Batch size 128 + - "128q1s512" + - "128q1s1k" + - "128q1s2k" + - "128q1s4k" + - "128q1s8k" + - "128q1s16k" + - "128q1s32k" + + # Batch size 256 + - "256q1s1k" + - "256q1s512" + - "256q1s2k" + - "256q1s4k" + - "256q1s8k" + - "256q1s16k" + - "256q1s32k" + +# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER +backends: + - FLASH_ATTN + - TRITON_ATTN + - FLASHINFER + +device: "cuda:0" +profile_memory: false diff --git a/benchmarks/attention_benchmarks/configs/standard_prefill.yaml b/benchmarks/attention_benchmarks/configs/standard_prefill.yaml new file mode 100644 index 00000000000..278b6347f65 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/standard_prefill.yaml @@ -0,0 +1,108 @@ +# Standard attention prefill benchmark configuration +# Sweeps num_q_heads and num_kv_heads to isolate effects of: +# 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads) +# 2. Absolute head count (fixed 4:1 ratio, vary scale) + +model: + num_layers: 32 + num_q_heads: 32 # Base value, overridden by sweep + num_kv_heads: 8 # Base value, overridden by sweep + head_dim: 128 + block_size: 16 + +# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and +# head_dim where it differs from the base (128). Head counts are per-GPU +# (i.e. after TP sharding). +# +# Group A — vary GQA ratio (fixed q=32, head_dim=128): +# 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA) +# +# Groups B-E — real model configs at various TP degrees: +# Model head_dim Full TP2 TP4 TP8 +# Llama 3 8B 128 32:8 16:4 8:2 4:1 +# Llama 3 70B 128 64:8 32:4 16:2 8:1 +# GPT-OSS 120B 64 64:8 32:4 16:2 8:1 +# Llama 3 405B 128 128:8 64:4 32:2 16:1 +model_parameter_sweep: + values: + # --- head_dim=128 (Llama 3 family) --- + - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1 + - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1 + - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8 + - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4 + - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2 + - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1 + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1 + - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8 + - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4 + - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2 + - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1 + # --- head_dim=64 (GPT-OSS 120B) --- + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1 + label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}" + +batch_specs: + # ---- batch_size x prefill_len grid (prefill: q_len == seq_len) ---- + # Total tokens = batch_size * prefill_len, and prefill compute scales with + # prefill_len^2, so the largest cells are expensive. Trim batch sizes or + # lengths for quick iteration. + + # Batch size 1 + - "q512" + - "q1k" + - "q2k" + - "q4k" + - "q8k" + - "q16k" + - "q32k" + + # Batch size 2 + - "2q512" + - "2q1k" + - "2q2k" + - "2q4k" + - "2q8k" + - "2q16k" + - "2q32k" + + # Batch size 4 + - "4q512" + - "4q1k" + - "4q2k" + - "4q4k" + - "4q8k" + - "4q16k" + - "4q32k" + + # Batch size 8 + - "8q512" + - "8q1k" + - "8q2k" + - "8q4k" + - "8q8k" + - "8q16k" + - "8q32k" + + # Batch size 16 + - "16q512" + - "16q1k" + - "16q2k" + - "16q4k" + - "16q8k" + - "16q16k" + - "16q32k" + +# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER +backends: + - FLASH_ATTN + - TRITON_ATTN + - FLASHINFER + +device: "cuda:0" +profile_memory: false diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index abab1e2edba..c9b3fb29bb9 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -8,6 +8,8 @@ This module provides helpers for running MLA backends without needing full VllmConfig integration. """ +import statistics + import numpy as np import torch from batch_spec import parse_batch_spec @@ -17,6 +19,8 @@ from common import ( MockIndexer, MockKVBProj, MockLayer, + run_do_bench, + run_ncu_profile, setup_mla_dims, ) @@ -704,6 +708,8 @@ def _run_single_benchmark( device: torch.device, indexer=None, kv_cache_dtype: str | None = None, + output_scale: float | None = None, + fuse_quant_op: bool = False, ) -> BenchmarkResult: """ Run a single benchmark iteration. @@ -717,6 +723,11 @@ def _run_single_benchmark( mla_dims: MLA dimension configuration device: Target device indexer: Optional MockIndexer for sparse backends + output_scale: Static per-tensor FP8 scale for prefill output. None + keeps the plain bf16 output (no quantization). + fuse_quant_op: With output_scale set, True lets the prefill kernel write + FP8 directly; False runs bf16 attention then a standalone static-FP8 + quant. The delta isolates the saved post-quant kernel. Returns: BenchmarkResult with timing statistics @@ -820,63 +831,86 @@ def _run_single_benchmark( num_prefill, mla_dims, query_fmt, device, torch.bfloat16 ) - # Build forward function + # Prefill FP8 output: fused (kernel writes e4m3) vs separate post-quant. + prefill_fp8_output = None + prefill_output_scale = None + prefill_quant_op = None + if has_prefill and output_scale is not None: + from vllm.platforms import current_platform + + prefill_output_scale = torch.tensor( + [output_scale], device=device, dtype=torch.float32 + ) + if fuse_quant_op: + prefill_fp8_output = torch.empty_like( + prefill_inputs["output"], dtype=current_platform.fp8_dtype() + ) + else: + from vllm.model_executor.layers.quantization.input_quant_fp8 import ( + QuantFP8, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + ) + + prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + + fused_output = output_scale is not None and fuse_quant_op + + # Build forward function (runs a single decode/prefill pass) def forward_fn(): results = [] if has_decode: results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)) if has_prefill: - results.append( - impl.forward_mha( - prefill_inputs["q"], - prefill_inputs["k_c_normed"], - prefill_inputs["k_pe"], - kv_cache, - metadata, - prefill_inputs["k_scale"], - prefill_inputs["output"], - ) + out = impl.forward_mha( + prefill_inputs["q"], + prefill_inputs["k_c_normed"], + prefill_inputs["k_pe"], + kv_cache, + metadata, + prefill_inputs["k_scale"], + prefill_fp8_output if fused_output else prefill_inputs["output"], + prefill_output_scale if fused_output else None, ) + if fused_output: + out = prefill_fp8_output + elif prefill_quant_op is not None: + out, _ = prefill_quant_op( + prefill_inputs["output"], prefill_output_scale + ) + results.append(out) return results[0] if len(results) == 1 else tuple(results) - # Warmup - for _ in range(config.warmup_iters): - forward_fn() - torch.accelerator.synchronize() - - # Optionally capture a CUDA graph after warmup. - # Graph replay eliminates CPU launch overhead so timings reflect pure - # kernel time. - if config.use_cuda_graphs: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - forward_fn() - benchmark_fn = graph.replay - else: - benchmark_fn = forward_fn - - # Benchmark - times = [] - for _ in range(config.repeats): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - start.record() + def benchmark_fn(): for _ in range(config.num_layers): - benchmark_fn() - end.record() + forward_fn() - torch.accelerator.synchronize() - elapsed_ms = start.elapsed_time(end) - times.append(elapsed_ms / 1000.0 / config.num_layers) + if config.ncu_profile: + run_ncu_profile(benchmark_fn) + return BenchmarkResult( + config=config, + mean_time=0.0, + median_time=0.0, + std_time=0.0, + min_time=0.0, + max_time=0.0, + throughput_tokens_per_sec=0.0, + ) + + all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms) + + # Convert ms to seconds per layer + times = [t / 1000.0 / config.num_layers for t in all_ms] + mean_time = statistics.mean(times) - mean_time = float(np.mean(times)) return BenchmarkResult( config=config, mean_time=mean_time, - std_time=float(np.std(times)), - min_time=float(np.min(times)), - max_time=float(np.max(times)), + median_time=statistics.median(times), + std_time=statistics.stdev(times) if len(times) > 1 else 0.0, + min_time=min(times), + max_time=max(times), throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0, ) @@ -886,6 +920,8 @@ def _run_mla_benchmark_batched( configs_with_params: list[tuple], # [(config, threshold, num_splits), ...] index_topk: int = 2048, prefill_backend: str | None = None, + output_scale: float | None = None, + fuse_quant_op: bool = False, ) -> list[BenchmarkResult]: """ Unified batched MLA benchmark runner for all backends. @@ -1025,6 +1061,8 @@ def _run_mla_benchmark_batched( device, indexer=indexer, kv_cache_dtype=kv_cache_dtype, + output_scale=output_scale, + fuse_quant_op=fuse_quant_op, ) results.append(result) @@ -1052,6 +1090,8 @@ def run_mla_benchmark( num_kv_splits: int | None = None, index_topk: int = 2048, prefill_backend: str | None = None, + output_scale: float | None = None, + fuse_quant_op: bool = False, ) -> BenchmarkResult | list[BenchmarkResult]: """ Unified MLA benchmark runner for all backends. @@ -1071,6 +1111,9 @@ def run_mla_benchmark( index_topk: Topk value for sparse MLA backends (default 2048) prefill_backend: Prefill backend name (e.g., "fa3", "fa4"). When set, forces the specified FlashAttention version for prefill. + output_scale: Static per-tensor FP8 scale for prefill output (None = bf16). + fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill + kernel vs a standalone post-quant kernel. See _run_single_benchmark. Returns: BenchmarkResult (single mode) or list of BenchmarkResult (batched mode) @@ -1095,7 +1138,12 @@ def run_mla_benchmark( # Use unified batched execution results = _run_mla_benchmark_batched( - backend, configs_with_params, index_topk, prefill_backend=prefill_backend + backend, + configs_with_params, + index_topk, + prefill_backend=prefill_backend, + output_scale=output_scale, + fuse_quant_op=fuse_quant_op, ) # Return single result or list based on input diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index aa636cd9cb5..8cd20dced17 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -9,13 +9,20 @@ This module provides helpers for running standard attention backends """ import logging +import statistics import types from contextlib import contextmanager -import numpy as np import torch from batch_spec import parse_batch_spec, reorder_for_flashinfer -from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale +from common import ( + BenchmarkConfig, + BenchmarkResult, + MockLayer, + get_attention_scale, + run_do_bench, + run_ncu_profile, +) from vllm.config import ( CacheConfig, @@ -208,6 +215,13 @@ def _create_backend_impl( scale = get_attention_scale(config.head_dim) + # Set v_head_dim for diff-headdim backends. Always reset (defaulting to + # head_dim) so a prior run's value doesn't leak into this one via the + # backend's class-level state. + if hasattr(backend_class, "set_head_size_v"): + v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim + backend_class.set_head_size_v(v_dim) + impl = backend_class.get_impl_cls()( num_heads=config.num_q_heads, head_size=config.head_dim, @@ -300,6 +314,7 @@ def _create_input_tensors( from vllm.platforms import current_platform q_dtype = current_platform.fp8_dtype() + v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim q_list = [ torch.randn( total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype @@ -313,9 +328,7 @@ def _create_input_tensors( for _ in range(config.num_layers) ] v_list = [ - torch.randn( - total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype - ) + torch.randn(total_q, config.num_kv_heads, v_dim, device=device, dtype=dtype) for _ in range(config.num_layers) ] return q_list, k_list, v_list @@ -389,14 +402,17 @@ def _run_single_benchmark( device: torch.device, dtype: torch.dtype, ) -> tuple: - """Run single benchmark iteration with warmup and timing loop.""" - total_q = q_list[0].shape[0] - out = torch.empty( - total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype - ) + """Run single benchmark using triton's do_bench_cudagraph/do_bench. - # Warmup - for _ in range(config.warmup_iters): + Returns: + (timing_stats, mem_stats) where timing_stats is a dict with + mean/std/min/max in seconds per layer. + """ + total_q = q_list[0].shape[0] + v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim + out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype) + + def benchmark_fn(): for i in range(config.num_layers): impl.forward( layer, @@ -407,52 +423,22 @@ def _run_single_benchmark( attn_metadata, output=out, ) - torch.accelerator.synchronize() - # Optionally capture a CUDA graph after warmup. - # Graph replay eliminates CPU launch overhead so timings reflect pure - # kernel time. - if config.use_cuda_graphs: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for i in range(config.num_layers): - impl.forward( - layer, - q_list[i], - k_list[i], - v_list[i], - cache_list[i], - attn_metadata, - output=out, - ) - benchmark_fn = graph.replay + if config.ncu_profile: + run_ncu_profile(benchmark_fn) + timing_stats = dict.fromkeys(("mean", "median", "std", "min", "max"), 0.0) else: + all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms) - def benchmark_fn(): - for i in range(config.num_layers): - impl.forward( - layer, - q_list[i], - k_list[i], - v_list[i], - cache_list[i], - attn_metadata, - output=out, - ) - - # Benchmark - times = [] - for _ in range(config.repeats): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - start.record() - benchmark_fn() - end.record() - - torch.accelerator.synchronize() - elapsed_ms = start.elapsed_time(end) - times.append(elapsed_ms / 1000.0 / config.num_layers) # seconds per layer + # Convert ms to seconds per layer + times = [t / 1000.0 / config.num_layers for t in all_ms] + timing_stats = { + "mean": statistics.mean(times), + "std": statistics.stdev(times) if len(times) > 1 else 0.0, + "min": min(times), + "max": max(times), + "median": statistics.median(times), + } mem_stats = {} if config.profile_memory: @@ -461,7 +447,7 @@ def _run_single_benchmark( "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2, } - return times, mem_stats + return timing_stats, mem_stats # ============================================================================ @@ -541,6 +527,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: common_attn_metadata=common_metadata, ) + # Override num_splits for split-K testing (FlashAttention only) + if config.num_splits is not None and hasattr( + attn_metadata, "max_num_splits" + ): + attn_metadata.max_num_splits = config.num_splits + # Only quantize queries when the impl supports it quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr( impl, "supports_quant_query_input", False @@ -553,7 +545,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: config, max_num_blocks, backend_class, device, dtype ) - times, mem_stats = _run_single_benchmark( + timing_stats, mem_stats = _run_single_benchmark( config, impl, layer, @@ -566,15 +558,16 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: dtype, ) - mean_time = np.mean(times) + mean_time = timing_stats["mean"] throughput = total_q / mean_time if mean_time > 0 else 0 return BenchmarkResult( config=config, mean_time=mean_time, - std_time=np.std(times), - min_time=np.min(times), - max_time=np.max(times), + median_time=timing_stats["median"], + std_time=timing_stats["std"], + min_time=timing_stats["min"], + max_time=timing_stats["max"], throughput_tokens_per_sec=throughput, memory_allocated_mb=mem_stats.get("allocated_mb"), memory_reserved_mb=mem_stats.get("reserved_mb"), diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py index a69637bfc43..6349095ad72 100644 --- a/benchmarks/backend_request_func.py +++ b/benchmarks/backend_request_func.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field import aiohttp import huggingface_hub.constants from tqdm.asyncio import tqdm -from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast +from transformers import AutoTokenizer, PythonBackend, TokenizersBackend # NOTE(simon): do not import vLLM here so the benchmark script # can run without vLLM installed. @@ -609,7 +609,7 @@ def get_tokenizer( tokenizer_mode: str = "auto", trust_remote_code: bool = False, **kwargs, -) -> PreTrainedTokenizer | PreTrainedTokenizerFast: +) -> PythonBackend | TokenizersBackend: if pretrained_model_name_or_path is not None and not os.path.exists( pretrained_model_name_or_path ): diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py index 6056fcdd072..f0a35a0cf15 100644 --- a/benchmarks/benchmark_hidden_state_extraction.py +++ b/benchmarks/benchmark_hidden_state_extraction.py @@ -92,7 +92,6 @@ def run_baseline( llm = LLM( model=model, enable_prefix_caching=False, - enable_chunked_prefill=False, **extra_args, ) sampling_params = SamplingParams(max_tokens=1) @@ -194,7 +193,6 @@ async def _run_extraction_async( engine_args = AsyncEngineArgs( model=model, enable_prefix_caching=False, - enable_chunked_prefill=False, max_num_batched_tokens=40960, max_model_len=40960, speculative_config={ diff --git a/benchmarks/benchmark_pin_memory.py b/benchmarks/benchmark_pin_memory.py new file mode 100644 index 00000000000..63a6b75d914 --- /dev/null +++ b/benchmarks/benchmark_pin_memory.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark and regression-test pinned (page-locked) CPU memory for vLLM. + +Verifies that enabling pinned memory does not regress throughput or latency +compared to unpinned memory. Each condition runs in an isolated ``spawn`` +subprocess so both start from a cold CUDA context, giving an unbiased +comparison. + +Usage +----- +Run all tests with the default model:: + + python benchmarks/benchmark_pin_memory.py -v + +Override the model and optional max-model-len:: + + python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B -v + python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B \ + --max-model-len 8192 -v + +Run only throughput or latency tests:: + + python benchmarks/benchmark_pin_memory.py -v -k test_throughput + python benchmarks/benchmark_pin_memory.py -v -k test_latency + +Run only the v1 or v2 runner variant:: + + python benchmarks/benchmark_pin_memory.py -v -k v1 + python benchmarks/benchmark_pin_memory.py -v -k v2 + +Note: on WSL2, v1 runner tests are skipped because pin memory is not available +for the v1 runner without cpu_offload_gb. Run on other platforms to exercise v1. +""" + +import argparse +import json +import multiprocessing +import sys +import tempfile + +import pytest + +# Allow up to 2% degradation. Both benchmark runs start from an identical +# cold CUDA context (separate spawn subprocesses), so the measured difference +# reflects the genuine pin_memory overhead rather than cold/warm ordering bias. +_THROUGHPUT_TOLERANCE = 0.98 +_THROUGHPUT_NUM_REQUESTS = 200 +_THROUGHPUT_INPUT_LEN = 128 +_THROUGHPUT_OUTPUT_LEN = 512 +_THROUGHPUT_MAX_NUM_SEQS = 128 + +# Latency benchmark constants — match latency.py defaults. +_LATENCY_TOLERANCE = 1.02 # Allow up to 2% latency regression. +_LATENCY_BATCH_SIZE = 64 +_LATENCY_INPUT_LEN = 32 +_LATENCY_OUTPUT_LEN = 128 +_LATENCY_WARMUP_ITERS = 5 +_LATENCY_BENCH_ITERS = 15 + +_DEFAULT_MODEL = "unsloth/Qwen3-1.7B" +_DEFAULT_MAX_MODEL_LEN = 16384 + + +def _benchmark_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--model", default=_DEFAULT_MODEL) + parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN) + args, _ = parser.parse_known_args() + return args + + +@pytest.fixture +def model() -> str: + return _benchmark_args().model + + +@pytest.fixture +def max_model_len() -> int: + return _benchmark_args().max_model_len + + +def _skip_if_pin_memory_not_available(engine_args_kwargs: dict) -> None: + """Skip the current pytest test if pin_memory is unavailable for this config.""" + import vllm.utils.platform_utils as pu + from vllm.config import set_current_vllm_config + from vllm.engine.arg_utils import EngineArgs + + vllm_config = EngineArgs(**engine_args_kwargs).create_engine_config() + with set_current_vllm_config(vllm_config): + pu.is_pin_memory_available.cache_clear() + if not pu.is_pin_memory_available(): + import os + + runner = "v2" if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1" else "v1" + model = engine_args_kwargs.get("model", "unknown") + print( + f"\033[33mSKIP: pin_memory not available for " + f"{runner} runner, model={model}\033[0m" + ) + pytest.skip("pin_memory not available for this configuration") + + +def _throughput_worker( + pin: bool, + engine_args_kwargs: dict, + q: "multiprocessing.Queue[float]", + v2_mode: bool = False, +) -> None: + """Run throughput benchmark in a fresh spawn subprocess. + + Delegates to vllm/benchmarks/throughput.py main() using the random dataset, + so the methodology matches the official benchmark. Results are written to a + temp JSON file and forwarded through the queue as tokens/s. + + v2_mode: when True, monkeypatches is_uva_available() to always return True + so the v2 model runner's UVA buffers remain functional even when pin=False. + This isolates the non-UVA pin_memory paths in v2. + """ + import vllm.utils.platform_utils as pu + from vllm.platforms import current_platform + + pu.is_pin_memory_available.cache_clear() + pu.is_uva_available.cache_clear() + type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin) + if v2_mode: + pu.is_uva_available = lambda: True + + from vllm.benchmarks.throughput import add_cli_args + from vllm.benchmarks.throughput import main as throughput_main + + parser = argparse.ArgumentParser() + add_cli_args(parser) + args = parser.parse_args([]) + + for key, val in engine_args_kwargs.items(): + setattr(args, key, val) + args.max_num_seqs = _THROUGHPUT_MAX_NUM_SEQS + args.dataset_name = "random" + args.input_len = _THROUGHPUT_INPUT_LEN + args.output_len = _THROUGHPUT_OUTPUT_LEN + # Nullify defaults that conflict with explicit input/output_len. + args.random_input_len = None + args.random_output_len = None + args.random_prefix_len = None + args.num_prompts = _THROUGHPUT_NUM_REQUESTS + args.seed = 0 + args.disable_detokenize = True + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + tmp_path = f.name + args.output_json = tmp_path + + throughput_main(args) + + with open(tmp_path) as f: + results = json.load(f) + q.put(results["tokens_per_second"]) + + +def _run_throughput_benchmark( + pin: bool, + engine_args_kwargs: dict, + v2_mode: bool = False, +) -> float: + ctx = multiprocessing.get_context("spawn") + q = ctx.Queue() + p = ctx.Process( + target=_throughput_worker, + args=(pin, engine_args_kwargs, q, v2_mode), + ) + p.start() + p.join() + if p.exitcode != 0: + raise RuntimeError( + f"Throughput benchmark subprocess (pin={pin}) exited with code {p.exitcode}" + ) + return q.get() + + +def _latency_worker( + pin: bool, + engine_args_kwargs: dict, + q: "multiprocessing.Queue[dict]", + v2_mode: bool = False, +) -> None: + """Run latency benchmark in a fresh spawn subprocess. + + Follows latency.py methodology: fixed batch of dummy token IDs, warmup + iterations to reach steady state, then timed iterations reduced to avg + and percentiles. Results are written to a temp JSON file by latency_main + and forwarded through the queue. + """ + import vllm.utils.platform_utils as pu + from vllm.platforms import current_platform + + pu.is_pin_memory_available.cache_clear() + pu.is_uva_available.cache_clear() + type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin) + if v2_mode: + pu.is_uva_available = lambda: True + + from vllm.benchmarks.latency import add_cli_args + from vllm.benchmarks.latency import main as latency_main + + parser = argparse.ArgumentParser() + add_cli_args(parser) + args = parser.parse_args([]) + + for key, val in engine_args_kwargs.items(): + setattr(args, key, val) + args.input_len = _LATENCY_INPUT_LEN + args.output_len = _LATENCY_OUTPUT_LEN + args.batch_size = _LATENCY_BATCH_SIZE + args.num_iters_warmup = _LATENCY_WARMUP_ITERS + args.num_iters = _LATENCY_BENCH_ITERS + args.profile = False + args.disable_detokenize = True + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + tmp_path = f.name + args.output_json = tmp_path + + latency_main(args) + + with open(tmp_path) as f: + results = json.load(f) + q.put(results) + + +def _run_latency_benchmark( + pin: bool, + engine_args_kwargs: dict, + v2_mode: bool = False, +) -> dict: + ctx = multiprocessing.get_context("spawn") + q = ctx.Queue() + p = ctx.Process( + target=_latency_worker, + args=(pin, engine_args_kwargs, q, v2_mode), + ) + p.start() + p.join() + if p.exitcode != 0: + raise RuntimeError( + f"Latency benchmark subprocess (pin={pin}) exited with code {p.exitcode}" + ) + return q.get() + + +@pytest.mark.parametrize( + "test_v2_runner", + [ + pytest.param(False, id="v1"), + pytest.param(True, id="v2"), + ], +) +class TestPinnedMemory: + """Verify pinned memory yields >= throughput vs unpinned via real vLLM inference.""" + + def test_throughput(self, monkeypatch, test_v2_runner, model, max_model_len): + """Benchmark throughput with pin_memory forced on then off. + + Delegates to vllm/benchmarks/throughput.py main() with the random + dataset. Each condition runs in an isolated spawn subprocess so both + start from a cold CUDA context, giving an unbiased comparison. + """ + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0") + + engine_args_kwargs = dict( + model=model, + gpu_memory_utilization=0.88, + max_model_len=max_model_len, + enable_prefix_caching=False, + ) + + _skip_if_pin_memory_not_available(engine_args_kwargs) + + unpinned_tps = _run_throughput_benchmark( + False, engine_args_kwargs, v2_mode=test_v2_runner + ) + pinned_tps = _run_throughput_benchmark( + True, engine_args_kwargs, v2_mode=test_v2_runner + ) + + pct_diff = (pinned_tps - unpinned_tps) / unpinned_tps * 100 + runner = "v2" if test_v2_runner else "v1" + print( + f"\n=== Throughput results ({runner} runner, {model}) ===" + f"\npin_memory=True: {pinned_tps:.1f} tok/s" + f"\npin_memory=False: {unpinned_tps:.1f} tok/s" + f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)" + ) + + assert pinned_tps >= unpinned_tps * _THROUGHPUT_TOLERANCE, ( + f"Pinned throughput ({pinned_tps:.1f} tok/s) fell more than " + f"{(1.0 - _THROUGHPUT_TOLERANCE) * 100:.1f}% below " + f"unpinned ({unpinned_tps:.1f} tok/s)." + ) + + def test_latency(self, monkeypatch, test_v2_runner, model, max_model_len): + """Benchmark per-batch latency with pin_memory forced on then off. + + Follows vllm/benchmarks/latency.py: fixed dummy-token batch, warmup + iterations to reach steady state, then timed iterations reduced to avg + and percentiles. Subprocesses run serially so each gets a cold CUDA + context without GPU memory pressure from the other run. + """ + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0") + + engine_args_kwargs = dict( + model=model, + gpu_memory_utilization=0.88, + max_model_len=max_model_len, + enable_prefix_caching=False, + ) + + _skip_if_pin_memory_not_available(engine_args_kwargs) + + unpinned = _run_latency_benchmark( + False, engine_args_kwargs, v2_mode=test_v2_runner + ) + pinned = _run_latency_benchmark( + True, engine_args_kwargs, v2_mode=test_v2_runner + ) + + pct_diff = ( + (pinned["avg_latency"] - unpinned["avg_latency"]) + / unpinned["avg_latency"] + * 100 + ) + runner = "v2" if test_v2_runner else "v1" + print( + f"\n=== Latency results ({runner} runner, {model}) ===" + f"\npin_memory=True: avg={pinned['avg_latency']:.3f}s" + f" p50={pinned['percentiles']['50']:.3f}s" + f" p99={pinned['percentiles']['99']:.3f}s" + f"\npin_memory=False: avg={unpinned['avg_latency']:.3f}s" + f" p50={unpinned['percentiles']['50']:.3f}s" + f" p99={unpinned['percentiles']['99']:.3f}s" + f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)" + ) + + assert pinned["avg_latency"] <= unpinned["avg_latency"] * _LATENCY_TOLERANCE, ( + f"Pinned avg latency ({pinned['avg_latency']:.3f}s) exceeded " + f"unpinned ({unpinned['avg_latency']:.3f}s) by more than " + f"{(_LATENCY_TOLERANCE - 1.0) * 100:.1f}%." + ) + + +if __name__ == "__main__": + _parser = argparse.ArgumentParser(add_help=False) + _parser.add_argument("--model", default=_DEFAULT_MODEL) + _parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN) + _, _remaining = _parser.parse_known_args() + sys.exit(pytest.main([__file__] + _remaining)) diff --git a/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py new file mode 100644 index 00000000000..09a01b301be --- /dev/null +++ b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2025 FlyDSL Project Contributors + +import json +import os + +import torch +from aiter.test_common import run_perftest + +from vllm.model_executor.layers.fused_moe import fused_experts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + compressed_tensors_moe_w4a16_flydsl, +) +from vllm.utils.platform_utils import get_device_name_as_file_name + +RoutingBuffers = tuple[ + torch.Tensor, # sorted_token_ids + torch.Tensor, # sorted_weights + torch.Tensor, # sorted_expert_ids + torch.Tensor, # num_valid_ids (shape [1], i32) + int, # sorted_size + int, # blocks +] + +MODEL_PARAMS_TO_TUNE = [ + # (num_experts, inter_dim, hidden_size, topk) + (384, 256, 7168, 8), # Kimi K2.5 TP=8 + (384, 512, 7168, 8), # Kimi K2.5 TP=4 +] + +NUM_TOKENS_TO_TUNE = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, +] + +TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] +TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] + +TILE_CONFIGS = [] +for tile_m in TILE_M_SEARCH_SPACE: + for tile_n in TILE_N_SEARCH_SPACE: + for tile_k in TILE_K_SEARCH_SPACE: + for tile_n2 in TILE_N2_SEARCH_SPACE: + for tile_k2 in TILE_K2_SEARCH_SPACE: + TILE_CONFIGS.append( + { + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "tile_n2": tile_n2, + "tile_k2": tile_k2, + } + ) + + +def tune_flydsl_moe_w4a16( + device: str = "cuda", num_iters: int = 100, num_warmup: int = 10 +): + packed_factor = 8 + w13_num_shards = 2 + params_dtype = torch.bfloat16 + group_size = 32 + scale_factor = 0.01 + + for model_params in MODEL_PARAMS_TO_TUNE: + num_experts = model_params[0] + inter_dim = model_params[1] + hidden_size = model_params[2] + topk = model_params[3] + print( + f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, " + f"hidden_size={hidden_size}, topk={topk}...\n" + ) + + w2_scales_size = inter_dim + num_groups_w2 = w2_scales_size // group_size + num_groups_w13 = hidden_size // group_size + + w13_weight = torch.randint( + 0, + 255, + (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim), + dtype=torch.int32, + device=device, + ) + + w2_weight = torch.randint( + 0, + 255, + (num_experts, inter_dim // packed_factor, hidden_size), + dtype=torch.int32, + device=device, + ) + w13_scale = scale_factor * torch.randn( + num_experts, + num_groups_w13, + w13_num_shards * inter_dim, + dtype=params_dtype, + device=device, + ) + w2_scale = scale_factor * torch.randn( + num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device + ) + + w13 = w13_weight + w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13) + w13 = w13.view(-1).contiguous() + + w2 = w2_weight + w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2) + w2 = w2.view(-1).contiguous() + + w13_scale_flydsl = w13_scale + w2_scale_flydsl = w2_scale + + if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: + E, G, N = w13_scale.shape + w13_scale_flydsl = ( + w13_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: + w13_scale_flydsl = w13_scale_flydsl.squeeze(1) + + if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: + E, G, N = w2_scale.shape + w2_scale_flydsl = ( + w2_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: + w2_scale_flydsl = w2_scale_flydsl.squeeze(1) + + w13_scale_flydsl = w13_scale_flydsl.contiguous() + w2_scale_flydsl = w2_scale_flydsl.contiguous() + + w13.is_shuffled = True + w2.is_shuffled = True + + w13_weight_scale = w13_scale.transpose(1, 2).contiguous() + w2_weight_scale = w2_scale.transpose(1, 2).contiguous() + w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8) + w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8) + + moe_quant_config = int4_w4a16_moe_quant_config( + w1_scale=w13_weight_scale, + w2_scale=w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + + tuned_config = {} + + for num_tokens in NUM_TOKENS_TO_TUNE: + score = torch.rand( + (num_tokens, num_experts), device=device, dtype=torch.float32 + ) + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) + x = torch.randn( + (num_tokens, hidden_size), dtype=torch.bfloat16, device=device + ) + us_best = float("inf") + for tile_config in TILE_CONFIGS: + try: + tile_m = tile_config["tile_m"] + tile_n = tile_config["tile_n"] + tile_k = tile_config["tile_k"] + tile_n2 = tile_config["tile_n2"] + tile_k2 = tile_config["tile_k2"] + + model_dim = x.shape[1] + assert model_dim % 64 == 0 + assert model_dim % tile_k == 0 + assert inter_dim % tile_n == 0 + assert model_dim % tile_n2 == 0 + assert inter_dim % tile_k2 == 0 + assert ((tile_m * tile_k2) % 256) == 0 + bytes_per_thread_x = (tile_m * tile_k2) // 256 + assert (bytes_per_thread_x % 4) == 0 + + out, _us = run_perftest( + fused_flydsl_moe, + x, + w13, + w2, + num_experts, + inter_dim, + topk_weights, + topk_ids, + num_iters=num_iters, + num_warmup=num_warmup, + w1_scale=w13_scale_flydsl, + w2_scale=w2_scale_flydsl, + topk=topk_weights.shape[-1], + group_size=group_size, + doweight_stage1=False, + scale_is_bf16=True, + config=tile_config, + ) + torch.accelerator.synchronize() + except Exception: + torch.accelerator.synchronize() + continue + else: + us = _us.item() + if us < us_best: + out_ref = fused_experts( + x, + w13_weight_packed, + w2_weight_packed, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + global_num_experts=num_experts, + expert_map=None, + quant_config=moe_quant_config, + ) + try: + assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1) + except Exception: + continue + else: + print( + f"For [num_tokens={num_tokens}, num_experts={num_experts}, " # noqa: E501 + f"inter_dim={inter_dim}] found new best " # noqa: E501 + f"config={tile_config}, us={us:0.3f}" + ) + us_best = us + tuned_config[str(num_tokens)] = tile_config + device_name = get_device_name_as_file_name() + tuned_config_file_name = ( + f"E={num_experts},N={inter_dim},device_name={device_name}," + f"dtype=int4_w4a16,backend=flydsl.json" + ) + tuner_dir_path = os.path.dirname(os.path.realpath(__file__)) + store_path = os.path.join(tuner_dir_path, tuned_config_file_name) + with open(store_path, "w") as f: + json.dump(tuned_config, f, indent=4) + print( + f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n" # noqa: E501 + ) + + +if __name__ == "__main__": + tune_flydsl_moe_w4a16(device="cuda") diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index 36cbd715f18..c999c16021b 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/benchmarks/kernels/benchmark_fused_collective.py @@ -80,13 +80,17 @@ _FI_MAX_SIZES = { 2: 64 * MiB, # 64MB 4: 64 * MiB, # 64MB 8: 64 * MiB, # 64MB + 16: 64 * MiB, # 64MB (multi-node) } # Global workspace tensors for FlashInfer (keyed by backend name) _FI_WORKSPACES: dict = {} -# Backends to benchmark -FLASHINFER_BACKENDS = ["trtllm", "mnnvl"] +# Backends to benchmark. trtllm is single-node only and can hang cross-node, so +# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl. +FLASHINFER_BACKENDS = [ + b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b +] def setup_flashinfer_workspace( @@ -995,7 +999,10 @@ def main(): rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - device = torch.device(f"cuda:{rank}") + # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per + # node) map to a valid local GPU; falls back to global rank single-node. + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + device = torch.device(f"cuda:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index f885b1e0952..1531cc96920 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -391,16 +391,19 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int config = dict(zip(keys, config_values)) configs.append(config) - # Remove configs that are not compatible with fp8 block quantization - # BLOCK_SIZE_K must be a multiple of block_k - # BLOCK_SIZE_N must be a multiple of block_n + # Drop configs incompatible with fp8 block quantization. A tile must align + # to the quant-block scale grid, i.e. tile and block must divide one + # another. The kernel indexes scales per element (offs_bn // group_n, + # k_start // group_k), so a tile narrower than the block (e.g. N=64 with + # block_n=128) is valid -- and often faster at small batch. An exact + # multiple was required before, which dropped those smaller tiles entirely. if block_quant_shape is not None and not use_fp16: block_n, block_k = block_quant_shape[0], block_quant_shape[1] for config in configs[:]: - if ( - config["BLOCK_SIZE_K"] % block_k != 0 - or config["BLOCK_SIZE_N"] % block_n != 0 - ): + bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"] + n_aligned = bn % block_n == 0 or block_n % bn == 0 + k_aligned = bk % block_k == 0 or block_k % bk == 0 + if not (n_aligned and k_aligned): configs.remove(config) return configs @@ -792,6 +795,12 @@ def get_model_params(config): topk = text_config.num_experts_per_tok intermediate_size = text_config.moe_intermediate_size hidden_size = text_config.hidden_size + elif architecture == "DiffusionGemmaForBlockDiffusion": + text_config = config.get_text_config() + E = text_config.num_experts + topk = text_config.top_k_experts + intermediate_size = text_config.moe_intermediate_size + hidden_size = text_config.hidden_size elif architecture == "HunYuanMoEV1ForCausalLM": E = config.num_experts topk = config.moe_topk[0] diff --git a/benchmarks/kernels/benchmark_paged_attention.py b/benchmarks/kernels/benchmark_paged_attention.py index b6a0b7ad8ca..f4249d6270a 100644 --- a/benchmarks/kernels/benchmark_paged_attention.py +++ b/benchmarks/kernels/benchmark_paged_attention.py @@ -19,13 +19,11 @@ from vllm.utils.torch_utils import ( logger = init_logger(__name__) NUM_BLOCKS = 128 * 1024 -PARTITION_SIZE = 512 PARTITION_SIZE_ROCM = 256 @torch.inference_mode() def main( - version: str, num_seqs: int, seq_len: int, num_query_heads: int, @@ -82,27 +80,20 @@ def main( # Prepare for the paged attention kernel. output = torch.empty_like(query) - if version == "v2": - if current_platform.is_rocm(): - global PARTITION_SIZE - if not args.custom_paged_attn and not current_platform.is_navi(): - PARTITION_SIZE = 1024 - else: - PARTITION_SIZE = PARTITION_SIZE_ROCM - num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE - tmp_output = torch.empty( - size=(num_seqs, num_query_heads, num_partitions, head_size), - dtype=output.dtype, - device=output.device, - ) - exp_sums = torch.empty( - size=(num_seqs, num_query_heads, num_partitions), - dtype=torch.float32, - device=output.device, - ) - max_logits = torch.empty_like(exp_sums) + num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM + tmp_output = torch.empty( + size=(num_seqs, num_query_heads, num_partitions, head_size), + dtype=output.dtype, + device=output.device, + ) + exp_sums = torch.empty( + size=(num_seqs, num_query_heads, num_partitions), + dtype=torch.float32, + device=output.device, + ) + max_logits = torch.empty_like(exp_sums) - def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float: + def run_benchmark(num_iters: int, profile: bool = False) -> float: torch.accelerator.synchronize() if profile: torch.cuda.cudart().cudaProfilerStart() @@ -112,67 +103,26 @@ def main( k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) for _ in range(num_iters): - if version == "v1": - ops.paged_attention_v1( - output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - elif version == "v2": - if not args.custom_paged_attn: - ops.paged_attention_v2( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - else: - ops.paged_attention_rocm( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - None, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - else: - raise ValueError(f"Invalid version: {version}") + ops.paged_attention_rocm( + output, + exp_sums, + max_logits, + tmp_output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + None, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + ) torch.accelerator.synchronize() end_time = time.perf_counter() @@ -182,7 +132,6 @@ def main( # Warmup. print("Warming up...") - run_benchmark = run_cuda_benchmark run_benchmark(num_iters=3, profile=False) # Benchmark. @@ -195,12 +144,13 @@ def main( if __name__ == "__main__": logger.warning( - "This script benchmarks the paged attention kernel. " + "This script benchmarks the ROCm paged attention kernel. " "By default this is no longer used in vLLM inference." ) + if not current_platform.is_rocm(): + raise RuntimeError("This benchmark requires the ROCm platform.") parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.") - parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2") parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--seq-len", type=int, default=4096) parser.add_argument("--num-query-heads", type=int, default=64) @@ -208,7 +158,7 @@ if __name__ == "__main__": parser.add_argument( "--head-size", type=int, - choices=[64, 80, 96, 112, 120, 128, 192, 256], + choices=[64, 128], default=128, ) parser.add_argument("--block-size", type=int, choices=[16, 32], default=16) @@ -224,11 +174,7 @@ if __name__ == "__main__": choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"], default="auto", help="Data type for kv cache storage. If 'auto', will use model " - "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. " - "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)", - ) - parser.add_argument( - "--custom-paged-attn", action="store_true", help="Use custom paged attention" + "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)", ) args = parser.parse_args() print(args) @@ -236,7 +182,6 @@ if __name__ == "__main__": if args.num_query_heads % args.num_kv_heads != 0: raise ValueError("num_query_heads must be divisible by num_kv_heads") main( - version=args.version, num_seqs=args.batch_size, seq_len=args.seq_len, num_query_heads=args.num_query_heads, diff --git a/benchmarks/kernels/benchmark_relu_squared.py b/benchmarks/kernels/benchmark_relu_squared.py new file mode 100644 index 00000000000..00550ca475b --- /dev/null +++ b/benchmarks/kernels/benchmark_relu_squared.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Benchmark ReLUSquaredActivation: custom CUDA kernel vs forward_native, both +# eager and under torch.compile (Inductor fuses relu+square into one kernel). + +import itertools + +import torch +import torch.nn.functional as F + +import vllm.model_executor.layers.activation # noqa: F401 +from vllm.benchmarks.lib.utils import default_vllm_config +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed + +# Capped so the largest tensor stays under 2**31 elements: the shared activation +# kernel computes the per-token pointer offset (blockIdx.x * d) in 32-bit, which +# overflows for tensors with >2**32 elements. Realistic token counts are well +# below this; the kernel-vs-native gap is already clear at these sizes. +batch_size_range = [1, 16, 128] +seq_len_range = [1, 16, 64, 1024] +intermediate_size = [3072, 9728, 12288] +configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size)) + + +@default_vllm_config() +def benchmark_relu_squared( + batch_size: int, + seq_len: int, + intermediate_size: int, + provider: str, + dtype: torch.dtype, +): + device = "cuda" + num_tokens = batch_size * seq_len + set_random_seed(42) + torch.set_default_device(device) + + x = torch.randn(num_tokens, intermediate_size, dtype=dtype, device=device) + out = torch.empty_like(x) + + def native(x: torch.Tensor) -> torch.Tensor: + return torch.square(F.relu(x)) + + # Verify the custom kernel matches the native implementation before timing. + ref = native(x) + torch.ops._C.relu_squared(out, x) + torch.testing.assert_close(out, ref) + + if provider == "custom": + # Custom CUDA kernel — single fused kernel. + fn = lambda: torch.ops._C.relu_squared(out, x) + elif provider == "native": + # forward_native, eager — relu and square as separate ops. + fn = lambda: native(x) + elif provider == "native_compiled": + # forward_native under torch.compile — Inductor fuses relu+square. + # This is the real production baseline (custom ops are off when + # Inductor is enabled), so it is the comparison reviewers care about. + compiled = torch.compile(native) + compiled(x) # warm up / trigger compilation before timing + fn = lambda: compiled(x) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=[0.5, 0.2, 0.8] + ) + return ms, max_ms, min_ms + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark ReLUSquaredActivation: custom kernel vs native." + ) + parser.add_argument( + "--dtype", + type=str, + choices=["half", "bfloat16", "float"], + default="bfloat16", + ) + args = parser.parse_args() + + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype] + + perf_report = triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len", "intermediate_size"], + x_vals=configs, + line_arg="provider", + line_vals=["custom", "native_compiled", "native"], + line_names=[ + "Custom Kernel", + "Native (torch.compile)", + "Native (eager)", + ], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="ms", + plot_name="relu_squared-eager-performance", + args={}, + ) + ) + + perf_report( + lambda batch_size, seq_len, intermediate_size, provider: benchmark_relu_squared( + batch_size, seq_len, intermediate_size, provider, dtype + ) + ).run(print_data=True) diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py index 36dce1b6388..590d4cfdc6d 100644 --- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py +++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.platform_utils import get_device_name_as_file_name mp.set_start_method("spawn", force=True) @@ -264,7 +265,7 @@ def save_configs( input_type="fp8", ) -> None: os.makedirs(save_path, exist_ok=True) - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() json_file_name = ( f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8," f"block_shape=[{block_n},{block_k}].json" diff --git a/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py b/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py index aff443083a5..f5a5ed1dc55 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py @@ -7,6 +7,7 @@ import time import numpy as np import torch +from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.torch_utils import set_random_seed @@ -14,17 +15,15 @@ from vllm.utils.torch_utils import set_random_seed try: from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight except (ImportError, AttributeError) as e: - print("ERROR: CPU fused MoE operations are not available on this platform.") - print("This benchmark requires x86 CPU with proper vLLM CPU extensions compiled.") - print( - "The cpu_fused_moe kernel is typically available on Linux x86_64 " - "with AVX2/AVX512." - ) print(f"Import error: {e}") sys.exit(1) # ISA selection following test_cpu_fused_moe.py pattern -ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"] +ISA_CHOICES = ["vec"] +if torch.cpu._is_amx_tile_supported(): + ISA_CHOICES.append("amx") +if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + ISA_CHOICES.append("neon") @torch.inference_mode() @@ -145,7 +144,7 @@ if __name__ == "__main__": "--isa", type=str, choices=ISA_CHOICES, - default=ISA_CHOICES[0], + default="vec", help=f"ISA to use (available: {ISA_CHOICES})", ) parser.add_argument("--seed", type=int, default=0) diff --git a/benchmarks/kv_cache_watermark.sh b/benchmarks/kv_cache_watermark.sh new file mode 100755 index 00000000000..258afa9fce1 --- /dev/null +++ b/benchmarks/kv_cache_watermark.sh @@ -0,0 +1,248 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Reproducible demonstration of the KV cache watermark (`--watermark`) for +# reducing preemption thrashing. +# +# The watermark is the fraction of total KV cache blocks the scheduler keeps +# free when admitting a waiting/preempted request into the running queue. +# +# Why this workload triggers thrashing: +# Requests are admitted based on the KV cache they need *at admission time*. +# With `--scheduler-reserve-full-isl` (default) the input length is reserved up +# front, but the *output* length is unknown and unreserved. A decode-heavy +# workload (output >> input) at high concurrency therefore over-admits while +# requests are short, then runs out of KV cache as they all grow during decode +# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills +# them later, and repeats. The watermark keeps a block of KV cache free so +# running requests can grow into it instead of triggering this churn. +# +# This script launches `vllm serve` under a deliberately KV-constrained config +# and a decode-heavy workload, sweeping the watermark across several values, and +# reports the preemption count (scraped from /metrics), throughput, and latency +# percentiles for each. It then plots the results. +# +# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens +# (+/- 20% variance), sized to run each config for ~5 minutes. +# +# Usage: +# benchmarks/kv_cache_watermark.sh +# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh +# +# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it). +set -euo pipefail + +# ---- Config (override via environment) ------------------------------------- +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} +TP=${TP:-1} +PORT=${PORT:-8000} +URL="http://127.0.0.1:${PORT}" +# Constrain the KV cache to a *near-critical* size: large enough that the engine +# can run stably, but small enough that greedy over-admission tips it into +# preemption thrashing. (Independent of GPU size, so the demo is reproducible.) +# At the default workload this fits ~1.5x the mean concurrent KV demand. +KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} +# Optional weight loader (e.g. fastsafetensors on the GCP cluster). +LOAD_FORMAT=${LOAD_FORMAT:-auto} +# Decode-heavy workload: moderate input, long output, with length variance. The +# long output means preempted requests have generated a lot before eviction, so +# resuming them re-prefills a long sequence (high recomputation cost). +INPUT_LEN=${INPUT_LEN:-1000} +OUTPUT_LEN=${OUTPUT_LEN:-5000} +RANGE_RATIO=${RANGE_RATIO:-0.2} +CONCURRENCY=${CONCURRENCY:-128} +# Enough prompts to keep each config saturated for ~5+ minutes. +NUM_PROMPTS=${NUM_PROMPTS:-450} +OUTDIR=${OUTDIR:-./watermark_bench_results} +# Watermark fractions compared. "label value" per line; value=0 disables it. +CONFIGS=${CONFIGS:-"off 0 +w0.02 0.02 +w0.05 0.05 +w0.10 0.10 +w0.15 0.15"} + +KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024)) +mkdir -p "$OUTDIR" + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +scrape_preemptions() { + # Sum the vllm:num_preemptions_total counter across engines. + python - "${URL}/metrics" <<'PY' +import sys, urllib.request +total = 0.0 +try: + body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace") + for line in body.splitlines(): + if line.startswith("vllm:num_preemptions_total"): + total += float(line.rsplit(" ", 1)[-1]) +except Exception as e: # noqa: BLE001 + print(f"scrape error: {e}", file=sys.stderr) +print(int(total)) +PY +} + +wait_for_server() { + for _ in $(seq 1 300); do + if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: server process exited during startup" >&2; return 1 + fi + sleep 5 + done + echo "ERROR: server did not become ready" >&2; return 1 +} + +run_one() { + local label=$1 watermark=$2 + echo + echo "==================== watermark: ${label} (${watermark}) ====================" + vllm serve "$MODEL" \ + --tensor-parallel-size "$TP" \ + --load-format "$LOAD_FORMAT" \ + --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \ + --max-model-len "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --no-enable-prefix-caching \ + --watermark "$watermark" \ + --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 & + SERVER_PID=$! + wait_for_server + sleep 5 + + local pre post + pre=$(scrape_preemptions) + vllm bench serve \ + --backend vllm \ + --base-url "$URL" \ + --model "$MODEL" \ + --dataset-name random \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --ignore-eos \ + --num-prompts "$NUM_PROMPTS" \ + --max-concurrency "$CONCURRENCY" \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --metric-percentiles "50,90,99" \ + --save-result \ + --result-dir "$OUTDIR" \ + --result-filename "bench_${label}.json" + post=$(scrape_preemptions) + echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt" + + kill "$SERVER_PID" 2>/dev/null || true + for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done + SERVER_PID="" + sleep 10 +} + +: >"${OUTDIR}/preemptions.txt" +while read -r label watermark; do + [[ -z "${label:-}" ]] && continue + run_one "$label" "$watermark" +done <<<"$CONFIGS" + +echo +echo "==================== summary ====================" +python - "$OUTDIR" <<'PY' +import json, os, sys +outdir = sys.argv[1] +pre = {} +order = [] +for line in open(os.path.join(outdir, "preemptions.txt")): + label, watermark, n = line.split() + pre[label] = (float(watermark), int(n)) + order.append(label) + +def g(d, *names): + for n in names: + if d.get(n) is not None: + return d[n] + return float("nan") + +cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s", + "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"] +print(" ".join(f"{c:>10}" for c in cols)) +rows = [] +for label in order: + watermark, n = pre[label] + d = json.load(open(os.path.join(outdir, f"bench_{label}.json"))) + rows.append(dict( + label=label, watermark=watermark, preempt=n, + out_tok_s=g(d, "output_throughput"), + req_s=g(d, "request_throughput"), + ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"), + ttft_p99=g(d, "p99_ttft_ms"), + itl_p99=g(d, "p99_itl_ms"), + e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"), + )) + print(" ".join(f"{str(v):>10}" for v in [ + label, watermark, n, + f"{rows[-1]['out_tok_s']:.0f}", + f"{rows[-1]['req_s']:.3f}", + f"{rows[-1]['ttft_p50']/1000:.2f}", + f"{rows[-1]['ttft_p99']/1000:.2f}", + f"{rows[-1]['itl_p99']:.2f}", + f"{rows[-1]['e2el_p50']/1000:.1f}", + ])) +print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)") + +# ---- Plot ------------------------------------------------------------------- +try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception as e: # noqa: BLE001 + print(f"\n(skip plot: matplotlib unavailable: {e})") + sys.exit(0) + +x = [r["watermark"] for r in rows] +xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows] +idx = list(range(len(rows))) + +fig, axes = plt.subplots(2, 2, figsize=(12, 8)) +fig.suptitle( + f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}", + fontsize=12, +) + +ax = axes[0][0] +ax.bar(idx, [r["preempt"] for r in rows], color="tab:red") +ax.set_title("Preemptions (lower is better)") +ax.set_ylabel("preemptions") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[0][1] +ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green") +ax.set_title("Output throughput (higher is better)") +ax.set_ylabel("tokens/s") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][0] +ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue") +ax.set_title("Inter-token latency p99 (lower is better)") +ax.set_ylabel("ITL p99 (ms)") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][1] +ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50") +ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99") +ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50") +ax.set_title("Latency (lower is better)") +ax.set_ylabel("seconds") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) +ax.legend() + +fig.tight_layout(rect=(0, 0, 1, 0.95)) +out_png = os.path.join(outdir, "watermark_results.png") +fig.savefig(out_png, dpi=120) +print(f"\nWrote plot: {out_png}") +PY diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index 750adf797ed..5a60d9c6688 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -65,6 +65,32 @@ class RequestArgs(NamedTuple): limit_min_tokens: int # Use negative value for no limit limit_max_tokens: int # Use negative value for no limit timeout_sec: int + send_conversation_id: bool + headers: dict[str, str] + + +def parse_custom_header(header: str) -> tuple[str, str]: + separators = (":", "=") + for separator in separators: + if separator in header: + key, value = header.split(separator, 1) + key = key.strip() + value = value.strip() + if key: + return key, value + break + raise argparse.ArgumentTypeError( + "Headers must be provided as 'Header-Name: value' or 'Header-Name=value'" + ) + + +def build_request_headers( + api_key: str | None, custom_headers: list[tuple[str, str]] | None +) -> dict[str, str]: + headers = dict(custom_headers or []) + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return headers class BenchmarkArgs(NamedTuple): @@ -218,12 +244,11 @@ async def send_request( max_tokens: int | None = None, timeout_sec: int = 120, conversation_id: str | None = None, + headers: dict[str, str] | None = None, ) -> ServerResponse: payload = { "model": model, "messages": messages, - "seed": 0, - "temperature": 0.0, } if conversation_id is not None: @@ -233,15 +258,17 @@ async def send_request( payload["stream"] = True payload["stream_options"] = {"include_usage": False} - if min_tokens is not None: - payload["min_tokens"] = min_tokens + # if min_tokens is not None: + # payload["min_tokens"] = min_tokens if max_tokens is not None: payload["max_tokens"] = max_tokens - headers = {"Content-Type": "application/json"} + request_headers = {"Content-Type": "application/json"} if conversation_id is not None: - headers["X-Session-ID"] = str(conversation_id) + request_headers["X-Session-ID"] = str(conversation_id) + if headers is not None: + request_headers.update(headers) # Calculate the timeout for the request if max_tokens is not None: @@ -267,7 +294,7 @@ async def send_request( most_recent_timestamp: int = start_time async with session.post( - url=chat_url, json=payload, headers=headers, timeout=timeout + url=chat_url, json=payload, headers=request_headers, timeout=timeout ) as response: http_status = HTTPStatus(response.status) if http_status == HTTPStatus.OK: @@ -319,6 +346,8 @@ async def send_request( latency = time.perf_counter_ns() - start_time if ttft is None: + if stream: + valid_response = False # The response was a single chunk ttft = latency @@ -425,7 +454,8 @@ async def send_turn( min_tokens, max_tokens, req_args.timeout_sec, - conversation_id=conv_id, + conversation_id=conv_id if req_args.send_conversation_id else None, + headers=req_args.headers, ) if response.valid is False: @@ -874,6 +904,7 @@ def get_client_config( # Arguments for API requests chat_url = f"{args.url}/v1/chat/completions" model_name = args.served_model_name if args.served_model_name else args.model + headers = build_request_headers(args.api_key, args.header) req_args = RequestArgs( chat_url=chat_url, @@ -882,6 +913,8 @@ def get_client_config( limit_min_tokens=args.limit_min_tokens, limit_max_tokens=args.limit_max_tokens, timeout_sec=args.request_timeout_sec, + send_conversation_id=args.send_conversation_id, + headers=headers, ) return client_args, req_args @@ -1247,19 +1280,19 @@ def process_statistics( ) -async def get_server_info(url: str) -> None: +async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None: logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}") async with aiohttp.ClientSession() as session: # Get server version (not mandatory, "version" endpoint may not exist) url_version = f"{url}/version" - async with session.get(url_version) as response: + async with session.get(url_version, headers=headers) as response: if HTTPStatus(response.status) == HTTPStatus.OK: text = await response.text() logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}") # Get available models url_models = f"{url}/v1/models" - async with session.get(url_models) as response: + async with session.get(url_models, headers=headers) as response: if HTTPStatus(response.status) == HTTPStatus.OK: text = await response.text() logger.info(f"{Color.BLUE}Models:{Color.RESET}") @@ -1325,6 +1358,22 @@ async def main() -> None: help="Base URL for the LLM API server", ) + parser.add_argument( + "--api-key", + type=str, + default=None, + help="API key to send as an Authorization bearer token", + ) + parser.add_argument( + "--header", + action="append", + type=parse_custom_header, + default=None, + metavar="KEY=VALUE", + help="Custom request header. Can be specified multiple times. " + "Accepts 'Header-Name: value' or 'Header-Name=value'.", + ) + parser.add_argument( "-p", "--num-clients", @@ -1439,6 +1488,22 @@ async def main() -> None: help="Disable stream/streaming mode (set 'stream' to False in the API request)", ) + parser.add_argument( + "--send-conversation-id", + default=False, + action="store_true", + help=( + "Inject a `conversation_id` field into each Chat Completions " + "payload. This is a non-standard OpenAI extension consumed by " + "vLLM's disaggregated multi-turn proxy " + "(examples/disaggregated/disaggregated_serving/" + "disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. " + "Leave disabled (default) when targeting strict " + "OpenAI-compatible endpoints; enable when benchmarking the " + "disaggregated proxy." + ), + ) + parser.add_argument( "-e", "--excel-output", @@ -1527,7 +1592,8 @@ async def main() -> None: args.model, trust_remote_code=args.trust_remote_code ) - await get_server_info(args.url) + headers = build_request_headers(args.api_key, args.header) + await get_server_info(args.url, headers=headers) # Load the input file (either conversations of configuration file) logger.info(f"Reading input file: {args.input_file}") diff --git a/build_rust.sh b/build_rust.sh index 98871ec8abc..1efc1ce39f1 100755 --- a/build_rust.sh +++ b/build_rust.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Build the vllm-rs Rust frontend binary and install it into the vllm package. +# Build vLLM Rust artifacts and install them into the vllm package. # Usage: ./build_rust.sh [--debug] # # By default builds in release mode. Pass --debug for faster compile times @@ -8,8 +8,6 @@ set -euo pipefail REPO_ROOT="$(cd "$(dirname "$0")" && pwd)" -RUST_DIR="$REPO_ROOT/rust" -TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}" # Read the required toolchain from rust-toolchain.toml. TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/') @@ -27,18 +25,9 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then fi if [[ "${1:-}" == "--debug" ]]; then - PROFILE_ARGS=() - PROFILE_DIR="debug" + PROFILE_ARG="--debug" else - PROFILE_ARGS=(--release) - PROFILE_DIR="release" + PROFILE_ARG="--release" fi -cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \ - --manifest-path "$RUST_DIR/Cargo.toml" \ - --bin vllm-rs \ - --features native-tls-vendored - -mkdir -p "$(dirname "$TARGET_PATH")" -cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH" -echo "Installed vllm-rs to $TARGET_PATH" +python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG" diff --git a/build_vllm_ppc64le.sh b/build_vllm_ppc64le.sh new file mode 100644 index 00000000000..3c0b74cc74d --- /dev/null +++ b/build_vllm_ppc64le.sh @@ -0,0 +1,241 @@ +#!/bin/bash +set -eoux pipefail + +######################################## +# Resolve repo root (IMPORTANT) +######################################## +REPO_ROOT="$(pwd)" + +cd "$REPO_ROOT" + +######################################## +# DevPI configuration +######################################## + +IBM_DEVPI_URL=${IBM_DEVPI_URL:-"https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/"} +RHOAI_INDEX_URL=${RHOAI_INDEX_URL:-"https://console.redhat.com/api/pypi/public-rhai/rhoai/3.4/cpu-ubi9/simple/"} + +######################################## +# wheel dir +######################################## + +WHEEL_DIR=${WHEEL_DIR:-"/tmp/wheels"} +mkdir -p "$WHEEL_DIR" + +######################################## +# Helpers +######################################## +try_install_from_devpi() { + local pkg=$1 + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --no-build-isolation \ + "${pkg}" +} + +######################################## +# Package Versions +######################################## +cd "$REPO_ROOT" +TORCH_VERSION=${TORCH_VERSION:-$(grep -E '^torch==.+==\s*"ppc64le"' requirements/cpu.txt | grep -Eo '\b[0-9\.]+\b' || true)} +TORCH_VERSION=${TORCH_VERSION:-2.11.0} + +TORCHVISION_VERSION=${TORCHVISION_VERSION:-0.26.0} +TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION:-${TORCH_VERSION}} + +export TORCH_VERSION +export TORCHVISION_VERSION +export TORCHAUDIO_VERSION +export OPENCV_VERSION=${OPENCV_VERSION:-4.13.0.92} +export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} + +######################################## +# install system dependencies +######################################## + +rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm || true + +microdnf install -y \ + python3.12 python3.12-devel python3.12-pip gcc \ + git jq gcc-toolset-14 gcc-toolset-14-libatomic-devel \ + automake libtool clang-devel openssl-devel \ + harfbuzz-devel kmod lcms2-devel libimagequant-devel libjpeg-turbo-devel \ + llvm15-devel libraqm-devel libtiff-devel libwebp-devel libxcb-devel \ + ninja-build openjpeg2-devel pkgconfig \ + tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-17.el9.ppc64le.rpm + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-devel-3.14.0-17.el9.ppc64le.rpm + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-compiler-3.14.0-17.el9.ppc64le.rpm + +######################################## +# Python 3.12 virtual environment +######################################## + +python3.12 -m venv /opt/vllm +source /opt/vllm/bin/activate + +export PATH=/opt/vllm/bin:$PATH + +python --version + +######################################## +# install build tools (stable uv) +######################################## + +pip install -U pip setuptools-rust +pip install uv +pip install "setuptools<70" build wheel cmake auditwheel +uv pip install "setuptools<70" cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation + +######################################## +# Rust +######################################## + +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +source /root/.cargo/env + +######################################## +# Compiler env +######################################## + +source /opt/rh/gcc-toolset-14/enable + +export PATH=/usr/lib64/llvm15/bin:$PATH +export LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config +export CMAKE_ARGS="-DPython3_EXECUTABLE=python" + +export MAX_JOBS=${MAX_JOBS:-$(nproc)} +export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 + +######################################## +# Install Packages From Devpi +######################################## +uv pip install numpy==2.3.5 pillow==12.2.0 --extra-index-url "$IBM_DEVPI_URL" +try_install_from_devpi "opencv-python-headless==${OPENCV_VERSION}" +try_install_from_devpi "torch==${TORCH_VERSION}" +try_install_from_devpi "torchvision==${TORCHVISION_VERSION}" + +######################################## +# torch audio +######################################## + +TEMP_BUILD_DIR=$(mktemp -d) +cd "${TEMP_BUILD_DIR}" +export BUILD_SOX=1 BUILD_KALDI=1 BUILD_RNNT=1 USE_FFMPEG=0 USE_ROCM=0 USE_CUDA=0 +export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 +git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} +cd audio +#patching +sed -i ' +s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; +s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; +s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; +' tools/setup_helpers/extension.py +MAX_JOBS=${MAX_JOBS:-$(nproc)} \ +BUILD_VERSION=${TORCHAUDIO_VERSION} \ +uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation +uv pip install "${WHEEL_DIR}"/torchaudio*.whl +cd "${REPO_ROOT}" +rm -rf "${TEMP_BUILD_DIR}" + +######################################## +# Xgrammar +######################################## +uv pip install \ + "scikit-build-core==0.11.6" \ + "pyproject-metadata<0.8" \ + pathspec \ + packaging \ + distro \ + "setuptools<70" \ + setuptools_scm \ + cmake \ + ninja \ + pybind11 \ + nanobind +uv pip install apache-tvm-ffi==0.1.12 \ + --no-build-isolation \ + --no-cache + +TEMP_BUILD_DIR=$(mktemp -d) + +pushd "${TEMP_BUILD_DIR}" + +export CFLAGS="-fno-lto -mcpu=power9" +export CXXFLAGS="-fno-lto -mcpu=power9" +export LDFLAGS="-fno-lto" +export PATH=/opt/vllm/bin:$PATH + +export Python_EXECUTABLE=/opt/vllm/bin/python3 +export Python3_EXECUTABLE=/opt/vllm/bin/python3 +export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 + +export Python_ROOT_DIR=/opt/vllm +export Python3_ROOT_DIR=/opt/vllm + +git clone \ + --recursive \ + https://github.com/mlc-ai/xgrammar \ + -b "v${XGRAMMAR_VERSION}" + +cd xgrammar + +cp cmake/config.cmake . +export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} + +uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + +uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + +popd + +rm -rf "${TEMP_BUILD_DIR}" +cd "${REPO_ROOT}" + +######################################## +# RHOAI Binary Downloads +######################################## +pip download \ + --index-url "${RHOAI_INDEX_URL}" \ + --only-binary=:all: \ + --no-deps \ + llvmlite==0.47.0 \ + -d "${WHEEL_DIR}" + +pip download \ + --index-url "${RHOAI_INDEX_URL}" \ + --only-binary=:all: \ + --no-deps \ + Numba==0.65.0 \ + -d "${WHEEL_DIR}" + +######################################## +# install built wheels +######################################## +uv pip install setuptools_scm maturin setuptools-rust ninja scikit-build-core pybind11 nanobind \ + --no-build-isolation +uv pip install "${WHEEL_DIR}"/*.whl + +######################################## +# install remaining deps +######################################## + +sed -i.bak -e 's/.*torch.*//g' pyproject.toml requirements/*.txt + +uv pip install "setuptools>=78.1.1" --no-build-isolation + +export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib64/pkgconfig:/usr/lib64/pkgconfig + +uv pip install -r requirements/common.txt \ + -r requirements/cpu.txt \ + -r requirements/build/cpu.txt --index-strategy unsafe-best-match diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 6f836ff5354..3aca9bcea91 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -15,6 +15,7 @@ endif() # set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86}) set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16}) +set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -24,7 +25,10 @@ set (ENABLE_NUMA TRUE) # Check the compile flags # if(MACOSX_FOUND) + # Apple clang needs -Xpreprocessor to enable OpenMP. No runtime link is + # needed: _C is a dynamic_lookup bundle and resolves libomp from torch. list(APPEND CXX_COMPILE_FLAGS + "-Xpreprocessor" "-fopenmp" "-DVLLM_CPU_EXTENSION") else() list(APPEND CXX_COMPILE_FLAGS @@ -107,6 +111,13 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin + # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1 + # overrides the detection result. + if (ENABLE_RVV_BF16) + set(RVV_BF16_FOUND ON) + message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable") + endif() endif() if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) @@ -166,11 +177,19 @@ elseif (S390_FOUND) "-mtune=native") elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "RISC-V detected") + if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0) + message(FATAL_ERROR + "VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'") + endif() # 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. + # by default; set -DVLLM_RVV_VLEN=0 to force scalar RISC-V build. + # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV. if(NOT DEFINED VLLM_RVV_VLEN) # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. - if(EXISTS /proc/cpuinfo) + # Skip when cross-compiling — /proc/cpuinfo describes the build host. + if(CMAKE_CROSSCOMPILING) + message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo") + elseif(EXISTS /proc/cpuinfo) file(READ /proc/cpuinfo _cpuinfo) set(_best 0) foreach(_n IN ITEMS 128 256 512 1024) @@ -178,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") set(_best ${_n}) endif() endforeach() + # Only VLEN=128 and VLEN=256 are supported by the RVV kernels. + if(_best GREATER 256) + message(WARNING + "Detected VLEN=${_best} but only 128/256 are supported; " + "clamping to 256") + set(_best 256) + endif() if(_best GREATER 0) set(VLLM_RVV_VLEN ${_best}) endif() @@ -187,10 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") 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)") + "Please specify VLEN explicitly via CMAKE_ARGS:\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)") endif() endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) @@ -202,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "BF16 extension detected") set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) elseif(RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") + 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)") @@ -219,7 +244,7 @@ endif() # Build oneDNN for GEMM kernels -if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) +if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND) # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64 # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "") @@ -322,7 +347,7 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER") set(ONEDNN_BUILD_GRAPH "OFF") set(ONEDNN_ENABLE_JIT_PROFILING "ON") - set(ONEDNN_ENABLE_ITT_TASKS "OFF") + set(ONEDNN_ENABLE_ITT_TASKS "ON") set(ONEDNN_ENABLE_MAX_CPU_ISA "ON") set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON") set(ONEDNN_VERBOSE "ON") @@ -420,6 +445,8 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) set(VLLM_EXT_SRC "csrc/cpu/shm.cpp" "csrc/cpu/activation_lut_bf16.cpp" + "csrc/cpu/cpu_tanhf_neon.hpp" + "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) endif() @@ -435,6 +462,12 @@ if(USE_ONEDNN) ${VLLM_EXT_SRC}) endif() +if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") + set(VLLM_EXT_SRC + "csrc/cpu/sgl-kernels/gemm_int4.cpp" + ${VLLM_EXT_SRC}) +endif() + if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_SGL "csrc/cpu/sgl-kernels/conv.cpp" diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 183c42dc795..38d218d00ac 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -8,43 +8,74 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR}) set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR}) endif() +# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download). +# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW +# disallows one-argument Populate(dep) after Declare; MakeAvailable would run +# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load). if(DEEPGEMM_SRC_DIR) - FetchContent_Declare( - deepgemm - SOURCE_DIR ${DEEPGEMM_SRC_DIR} - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) + # cmake_path(ABSOLUTE_PATH ...) reads the path from ; NORMALIZE is a + # flag (no trailing path argument). Resolve relative paths against vLLM root. + set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}") + cmake_path(ABSOLUTE_PATH _deepgemm_user_src + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}" + NORMALIZE) + set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}") + if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}") + message(FATAL_ERROR + "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'") + endif() + set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}") + message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") 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 "" - ) + # Keep in sync with tools/install_deepgemm.sh + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") + # NOTE: This is currently targeting nv-dev branch due to sm120 support + set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a") + + set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") + if(NOT _deepgemm_fc_root) + set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps") + endif() + set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src") + set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build") + set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild") + + if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp") + set(deepgemm_SOURCE_DIR "${_deepgemm_src}") + set(deepgemm_BINARY_DIR "${_deepgemm_bin}") + else() + FetchContent_Populate( + deepgemm + SUBBUILD_DIR "${_deepgemm_sub}" + SOURCE_DIR "${_deepgemm_src}" + BINARY_DIR "${_deepgemm_bin}" + GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" + GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" + GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_PROGRESS TRUE + ) + endif() + message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}") 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 +# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream), +# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x +# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants. 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") +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") + else() + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f") + else() + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a") + endif() endif() cuda_archs_loose_intersection(DEEPGEMM_ARCHS diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 65986df5501..50ddc667a84 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 a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b + GIT_TAG b70aff3d110a2b1a037e62eac295166b5143643a GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake new file mode 100644 index 00000000000..052966b2755 --- /dev/null +++ b/cmake/external_projects/fmha_sm100.cmake @@ -0,0 +1,73 @@ +include(FetchContent) + +# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory +# instead of downloading. This is useful for local MSA development. +if(DEFINED ENV{FMHA_SM100_SRC_DIR}) + set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR}) +endif() + +if(FMHA_SM100_SRC_DIR) + FetchContent_Declare( + fmha_sm100 + SOURCE_DIR ${FMHA_SM100_SRC_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +else() + FetchContent_Declare( + fmha_sm100 + GIT_REPOSITORY https://github.com/vllm-project/MSA.git + GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73 + GIT_PROGRESS TRUE + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +endif() + +FetchContent_GetProperties(fmha_sm100) +if(NOT fmha_sm100_POPULATED) + FetchContent_Populate(fmha_sm100) +endif() +message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}") + +add_custom_target(fmha_sm100) + +set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100") + +install(FILES + "${FMHA_SM100_PY_ROOT}/__init__.py" + "${FMHA_SM100_PY_ROOT}/api.py" + "${FMHA_SM100_PY_ROOT}/bench_utils.py" + "${FMHA_SM100_PY_ROOT}/jit.py" + "${FMHA_SM100_PY_ROOT}/sparse.py" + "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py" + DESTINATION vllm/third_party/fmha_sm100 + COMPONENT fmha_sm100) + +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/" + DESTINATION vllm/third_party/fmha_sm100/csrc + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE + PATTERN ".git*" EXCLUDE) + +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/" + DESTINATION vllm/third_party/fmha_sm100/cute + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE + PATTERN ".git*" EXCLUDE) + +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/" + DESTINATION vllm/third_party/fmha_sm100/cutlass/include + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE + PATTERN ".git*" EXCLUDE) + +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/" + DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE + PATTERN ".git*" EXCLUDE) diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index 273fe754bed..29c5c6528b9 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR}) set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR}) endif() +# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid. +# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees. if(QUTLASS_SRC_DIR) - FetchContent_Declare( - qutlass - SOURCE_DIR ${QUTLASS_SRC_DIR} - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) + set(_qutlass_user_src "${QUTLASS_SRC_DIR}") + cmake_path(ABSOLUTE_PATH _qutlass_user_src + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}" + NORMALIZE) + set(QUTLASS_SRC_DIR "${_qutlass_user_src}") + if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}") + message(FATAL_ERROR + "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'") + endif() + set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}") + set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused") else() - FetchContent_Declare( - qutlass - GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git - GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65 - GIT_PROGRESS TRUE - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) -endif() + set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git") + set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65") -FetchContent_Populate(qutlass) + set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}") + if(NOT _qutlass_fc_root) + set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps") + endif() + set(_qutlass_src "${_qutlass_fc_root}/qutlass-src") + set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build") + set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild") + + if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp") + set(qutlass_SOURCE_DIR "${_qutlass_src}") + set(qutlass_BINARY_DIR "${_qutlass_bin}") + else() + FetchContent_Populate( + qutlass + SUBBUILD_DIR "${_qutlass_sub}" + SOURCE_DIR "${_qutlass_src}" + BINARY_DIR "${_qutlass_bin}" + GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}" + GIT_TAG "${_QUTLASS_UPSTREAM_TAG}" + GIT_PROGRESS TRUE + ) + endif() +endif() if(NOT qutlass_SOURCE_DIR) message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.") @@ -32,22 +54,35 @@ endif() message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}") +endif() + +# QUTLASS uses TARGET_CUDA_ARCH as a single preprocessor selector for all its +# sources. Do not compile a mixed SM100/SM120 arch list with one selector; prefer +# SM100 when both families are requested because that is the primary deployed +# target for this extension today. +if(QUTLASS_SM100_ARCHS) + set(QUTLASS_ARCHS "${QUTLASS_SM100_ARCHS}") + set(QUTLASS_TARGET_CC 100) + if(QUTLASS_SM120_ARCHS) + message(WARNING + "[QUTLASS] Both SM100 and SM120 archs were requested; selecting SM100 " + "because TARGET_CUDA_ARCH is a single compile-time selector.") + endif() +elseif(QUTLASS_SM120_ARCHS) + set(QUTLASS_ARCHS "${QUTLASS_SM120_ARCHS}") + set(QUTLASS_TARGET_CC 120) +else() + set(QUTLASS_ARCHS) endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) - - if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)") - set(QUTLASS_TARGET_CC 100) - elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?") - set(QUTLASS_TARGET_CC 120) - else() - message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.") - endif() - set(QUTLASS_SOURCES + csrc/qutlass_registration.cpp ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu @@ -66,8 +101,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h") list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}") + if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND + EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp") + list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}") + else() + get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY) + if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp") + list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include") + endif() + endif() elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h") - list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include") + list(APPEND QUTLASS_INCLUDES + "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include" + "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include") message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).") else() message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. " @@ -79,12 +125,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) CUDA_ARCHS "${QUTLASS_ARCHS}" ) - target_sources(_C PRIVATE ${QUTLASS_SOURCES}) - target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES}) - target_compile_definitions(_C PRIVATE + # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION. + # Keep it as its own extension (registers torch.ops._qutlass_C). + define_extension_target( + _qutlass_C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${QUTLASS_SOURCES} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + target_compile_definitions(_qutlass_C PRIVATE QUTLASS_DISABLE_PYBIND=1 TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC} - ) + CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS $<$:--expt-relaxed-constexpr --use_fast_math -O3> @@ -99,4 +156,5 @@ else() "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in " "CUDA_ARCHS='${CUDA_ARCHS}'.") endif() + add_custom_target(_qutlass_C) endif() diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 1e4feb0ff9e..97a8cfe87b7 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 dd62dac706b1cf7895bd99b18c6cb7e7e117ee25 + GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/cmake/utils.cmake b/cmake/utils.cmake index dd2034c1c5e..e3e766541df 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -487,9 +487,9 @@ endfunction() function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}") + cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") else() - cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}") + cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}") endif() set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE) endfunction() diff --git a/csrc/cpu/activation.cpp b/csrc/cpu/activation.cpp index 039b8d5c30d..2f06813a194 100644 --- a/csrc/cpu/activation.cpp +++ b/csrc/cpu/activation.cpp @@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out, // [..., d] }); } +void gelu_tanh(torch::Tensor& out, torch::Tensor& input) { + int num_tokens = input.numel() / input.size(-1); + int d = input.size(-1); + + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] { + CPU_KERNEL_GUARD_IN(gelu_tanh_impl) + activation_kernel( + num_tokens, d, input.data_ptr(), out.data_ptr()); + CPU_KERNEL_GUARD_OUT(gelu_tanh_impl) + }); +} + void gelu_new(torch::Tensor& out, torch::Tensor& input) { int num_tokens = input.numel() / input.size(-1); int d = input.size(-1); diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 26b881f4f14..fa22861157e 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -11,13 +11,26 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( return cpu_attention::Fp8KVCacheDataType::kAuto; } +bool cpu_attn_has_isa(const std::string& isa) { + if (isa == "rvv") { +#if defined(__riscv) && defined(__riscv_v_min_vlen) && \ + (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) + return true; +#else + return false; +#endif + } + return false; +} + 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, const torch::Tensor& seq_lens, at::ScalarType dtype, - const torch::Tensor& query_start_loc, const bool casual, + const torch::Tensor& query_start_loc, const bool causal, const int64_t window_size, const std::string& isa_hint, - const bool enable_kv_split) { + const bool enable_kv_split, + const std::optional& dynamic_causal) { cpu_attention::ISA isa; if (isa_hint == "amx") { isa = cpu_attention::ISA::AMX; @@ -44,24 +57,13 @@ torch::Tensor get_scheduler_metadata( input.head_dim = head_dim; input.query_start_loc = query_start_loc.data_ptr(); input.seq_lens = seq_lens.data_ptr(); - if (window_size != -1) { - input.left_sliding_window_size = window_size - 1; - if (casual) { - input.right_sliding_window_size = 0; - } else { - input.right_sliding_window_size = window_size - 1; - } - } else { - input.left_sliding_window_size = -1; - if (casual) { - input.right_sliding_window_size = 0; - } else { - input.right_sliding_window_size = -1; - } - } - input.casual = casual; + + input.sliding_window_size = window_size; + input.causal = causal; input.isa = isa; input.enable_kv_split = enable_kv_split; + input.dynamic_causal = + dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() { @@ -175,10 +177,11 @@ void cpu_attention_with_kv_cache( const torch::Tensor& seq_lens, // [num_tokens] const double scale, const bool causal, const std::optional& alibi_slopes, // [num_heads] - const int64_t sliding_window_left, const int64_t sliding_window_right, + const int64_t sliding_window, const torch::Tensor& block_table, // [num_tokens, max_block_num] const double softcap, const torch::Tensor& scheduler_metadata, - const std::optional& s_aux, // [num_heads] + const std::optional& s_aux, // [num_heads] + const std::optional& dynamic_causal, // [num_reqs] const double k_scale = 1.0, const double v_scale = 1.0, const std::string& kv_cache_dtype = "auto") { TORCH_CHECK_EQ(query.dim(), 3); @@ -220,13 +223,11 @@ void cpu_attention_with_kv_cache( input.alibi_slopes = alibi_slopes.has_value() ? alibi_slopes->data_ptr() : nullptr; input.s_aux = s_aux.has_value() ? s_aux->data_ptr() : nullptr; + input.dynamic_causal = + dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; input.scale = scale; input.causal = causal; - input.sliding_window_left = sliding_window_left; - input.sliding_window_right = sliding_window_right; - if (input.causal) { - input.sliding_window_right = 0; - } + input.sliding_window_size = sliding_window; input.softcap = static_cast(softcap); if (is_fp8) { diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 70081b36ee5..907f8895682 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -124,7 +124,7 @@ struct AttentionMetadata { workitem_group_num(workitem_group_num), reduction_item_num(reduction_item_num), reduction_split_num(reduction_split_num), - thread_num(omp_get_max_threads()), + thread_num(cpu_utils::get_max_threads()), effective_thread_num(thread_num), split_kv_q_token_num_threshold(split_kv_q_token_num_threshold), attention_scratchpad_size_per_thread(0), @@ -388,13 +388,13 @@ class AttentionScheduler { int32_t head_dim; int32_t* query_start_loc; int32_t* seq_lens; - int32_t left_sliding_window_size; - int32_t right_sliding_window_size; - bool casual; + int32_t sliding_window_size; + bool causal; cpu_attention::ISA isa; int32_t max_num_q_per_iter; // max Q head num can be hold in registers int32_t kv_block_alignment; // context length alignment requirement bool enable_kv_split; + bool* dynamic_causal; }; static constexpr int32_t MaxQTileIterNum = 128; @@ -403,8 +403,9 @@ class AttentionScheduler { : available_cache_size_(cpu_utils::get_available_l2_size()) {} torch::Tensor schedule(const ScheduleInput& input) const { - const bool casual = input.casual; - const int32_t thread_num = omp_get_max_threads(); + const bool causal = input.causal; + const bool is_dynamic_causal = input.dynamic_causal != nullptr; + const int32_t thread_num = cpu_utils::get_max_threads(); const int64_t cache_size = cpu_utils::get_available_l2_size(); const int32_t max_num_q_per_iter = input.max_num_q_per_iter; const int32_t kv_len_alignment = input.kv_block_alignment; @@ -416,8 +417,10 @@ class AttentionScheduler { has_decode_request = has_decode_request || (q_token_num == 1); decode_only_batch = decode_only_batch && (q_token_num == 1); } - int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv; - const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter; + const int32_t original_q_head_per_kv = + input.num_heads_q / input.num_heads_kv; + int32_t q_head_per_kv = original_q_head_per_kv; + const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter; const bool use_gqa_fast_path = supports_gqa && decode_only_batch; const bool use_gqa_scratchpad = supports_gqa && has_decode_request; if (!use_gqa_scratchpad) { @@ -434,8 +437,7 @@ class AttentionScheduler { const int32_t default_tile_token_num = default_tile_size / q_head_per_kv; const int32_t split_kv_q_token_num_threshold = input.enable_kv_split ? 1 : 0; - const int32_t left_sliding_window_size = input.left_sliding_window_size; - const int32_t right_sliding_window_size = input.right_sliding_window_size; + const int32_t sliding_window_size = input.sliding_window_size; TORCH_CHECK_LE(split_kv_q_token_num_threshold * q_head_per_kv, 16); // get total kv len @@ -444,7 +446,9 @@ class AttentionScheduler { const int32_t seq_len = input.seq_lens[req_id]; const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; - const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0); + const bool req_causal = + is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; @@ -456,7 +460,7 @@ class AttentionScheduler { const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num; const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos( kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right, - left_sliding_window_size, right_sliding_window_size); + sliding_window_size, req_causal); const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] = align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right, kv_len_alignment); @@ -484,7 +488,9 @@ class AttentionScheduler { const int32_t seq_len = input.seq_lens[req_id]; const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; - const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0); + const bool req_causal = + is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; int32_t local_split_id = 0; @@ -498,7 +504,7 @@ class AttentionScheduler { const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num; const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos( kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right, - left_sliding_window_size, right_sliding_window_size); + sliding_window_size, req_causal); const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] = align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right, kv_len_alignment); @@ -667,22 +673,62 @@ class AttentionScheduler { metadata_ptr->effective_thread_num = effective_thread_num; { - // when q_tile_size = max_num_q_per_iter, requires max - // attention_scratchpad_size AttentionScratchPad sc(0, *metadata_ptr, 0x0); - int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q( - cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size, - input.logits_buffer_elem_size, input.output_buffer_elem_size, - max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true); - sc.update(input.head_dim, input.q_buffer_elem_size, - input.logits_buffer_elem_size, input.output_buffer_elem_size, - max_num_q_per_iter, max_num_q_per_iter, n); + int64_t max_attention_scratchpad_size = 0; + + for (const AttentionWorkItemGroup& item : workitems) { + const bool curr_use_gqa = + use_gqa_fast_path || (supports_gqa && item.q_token_num == 1); + const int32_t curr_q_heads_per_kv = + curr_use_gqa ? original_q_head_per_kv : 1; + const int32_t curr_default_q_tile_token_num = + default_tile_size / curr_q_heads_per_kv; + + for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num; + q_token_offset += curr_default_q_tile_token_num) { + const int32_t actual_q_token_num = std::min( + curr_default_q_tile_token_num, item.q_token_num - q_token_offset); + const int32_t q_head_tile_size = + actual_q_token_num * curr_q_heads_per_kv; + const int32_t rounded_q_head_tile_size = + ((q_head_tile_size + max_num_q_per_iter - 1) / + max_num_q_per_iter) * + max_num_q_per_iter; + + const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q( + cache_size, input.head_dim, input.elem_size, + input.q_buffer_elem_size, input.logits_buffer_elem_size, + input.output_buffer_elem_size, max_num_q_per_iter, + kv_len_alignment, rounded_q_head_tile_size, + rounded_q_head_tile_size <= max_num_q_per_iter); + + sc.update(input.head_dim, input.q_buffer_elem_size, + input.logits_buffer_elem_size, + input.output_buffer_elem_size, max_num_q_per_iter, + rounded_q_head_tile_size, n); + + max_attention_scratchpad_size = std::max( + max_attention_scratchpad_size, sc.get_thread_scratchpad_size()); + } + } + metadata_ptr->attention_scratchpad_size_per_thread = - ((sc.get_thread_scratchpad_size() + 63) / 64) * 64; + ((max_attention_scratchpad_size + 63) / 64) * 64; + + int32_t max_reduction_q_head_tile_size = 0; + for (const ReductionWorkItemGroup& item : reduce_workitems) { + const bool curr_use_gqa = + use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1); + const int32_t curr_q_heads_per_kv = + curr_use_gqa ? original_q_head_per_kv : 1; + + max_reduction_q_head_tile_size = + std::max(max_reduction_q_head_tile_size, + item.q_token_id_num * curr_q_heads_per_kv); + } sc.update(0, metadata_ptr->reduction_split_num, input.head_dim, - q_head_per_kv * split_kv_q_token_num_threshold, - input.output_buffer_elem_size); + max_reduction_q_head_tile_size, input.output_buffer_elem_size); metadata_ptr->reduction_scratchpad_size_per_kv_head = ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64; } @@ -708,15 +754,41 @@ class AttentionScheduler { return metadata_tensor; } + FORCE_INLINE static std::pair calcu_sliding_window_size( + int32_t window_size, bool causal) { + int32_t left_sliding_window_size, right_sliding_window_size; + if (window_size != -1) { + left_sliding_window_size = window_size - 1; + if (causal) { + right_sliding_window_size = 0; + } else { + right_sliding_window_size = window_size - 1; + } + } else { + left_sliding_window_size = -1; + if (causal) { + right_sliding_window_size = 0; + } else { + right_sliding_window_size = -1; + } + } + + return {left_sliding_window_size, right_sliding_window_size}; + } + FORCE_INLINE static std::pair calcu_kv_tile_pos( int32_t kv_left_pos, int32_t kv_right_pos, int32_t q_left_pos, - int32_t q_right_pos, int32_t sliding_window_left, - int32_t sliding_window_right) { - if (sliding_window_left != -1) { - kv_left_pos = std::max(kv_left_pos, q_left_pos - sliding_window_left); + int32_t q_right_pos, int32_t window_size, bool causal) { + auto [left_sliding_window_size, right_sliding_window_size] = + calcu_sliding_window_size(window_size, causal); + + if (left_sliding_window_size != -1) { + kv_left_pos = + std::max(kv_left_pos, q_left_pos - left_sliding_window_size); } - if (sliding_window_right != -1) { - kv_right_pos = std::min(kv_right_pos, q_right_pos + sliding_window_right); + if (right_sliding_window_size != -1) { + kv_right_pos = + std::min(kv_right_pos, q_right_pos + right_sliding_window_size); } return {kv_left_pos, kv_right_pos}; } @@ -805,10 +877,10 @@ struct AttentionInput { int32_t* block_table; float* alibi_slopes; c10::BFloat16* s_aux; + bool* dynamic_causal; float scale; bool causal; - int32_t sliding_window_left; - int32_t sliding_window_right; + int32_t sliding_window_size; float softcap; // FP8 KV cache scales (used by FP8 attention implementations) float k_scale_fp8 = 1.0f; @@ -822,8 +894,8 @@ struct AttentionInput { logits_buffer_t *__restrict__ logits_buffer, \ float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \ float *__restrict__ sum_buffer, int32_t *__restrict__ block_table, \ - const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos, \ - const int32_t kv_tile_token_num, \ + const int32_t kv_end_pos, const int32_t kv_tile_start_pos, \ + const int32_t kv_tile_end_pos, const int32_t kv_tile_token_num, \ const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num, \ const int32_t q_token_num, const int32_t q_tile_start_pos, \ const int32_t q_heads_per_kv, const int32_t block_size, \ @@ -834,7 +906,7 @@ struct AttentionInput { #define CPU_ATTENTION_PARAMS \ q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer, \ - partial_q_buffer, max_buffer, sum_buffer, block_table, \ + partial_q_buffer, max_buffer, sum_buffer, block_table, kv_end_pos, \ kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num, \ kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos, \ q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \ @@ -857,12 +929,10 @@ struct VecTypeTrait { using vec_t = vec_op::BF16Vec16; }; -#if !defined(__powerpc__) template <> struct VecTypeTrait { using vec_t = vec_op::FP16Vec16; }; -#endif template void print_logits(const char* name, T* ptr, int32_t row, int32_t col, @@ -917,6 +987,7 @@ class AttentionMainLoop { // - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits // - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp // - block_table + // - kv_end_pos: un-aligned end position of KV cache // - kv_tile_start_pos: start position of KV cache, aligned to // BlockSizeAlignment // - kv_tile_end_pos: end position of KV cache, aligned to @@ -1043,7 +1114,7 @@ class AttentionMainLoop { } apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos, - kv_tile_start_pos, kv_tile_end_pos, q_token_num, + kv_end_pos, kv_tile_start_pos, kv_tile_end_pos, q_token_num, q_heads_per_kv, left_window_size, right_window_size); // if (debug_info){ @@ -1126,7 +1197,7 @@ class AttentionMainLoop { void apply_mask(logits_buffer_t* __restrict__ logits_buffer, const int64_t logits_buffer_stride, - const int32_t q_tile_start_pos, + const int32_t q_tile_start_pos, const int32_t kv_end_pos, const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos, const int32_t q_token_num, const int32_t q_heads_per_kv, @@ -1154,7 +1225,7 @@ class AttentionMainLoop { std::max(kv_tile_start_pos, curr_token_pos + sliding_window_right + 1)); } - return pos; + return std::min(pos, kv_end_pos); }(); int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos; @@ -1394,7 +1465,7 @@ class AttentionMainLoop { public: void operator()(const AttentionInput* input) { - const int thread_num = omp_get_max_threads(); + const int thread_num = cpu_utils::get_max_threads(); TORCH_CHECK_EQ(input->metadata->thread_num, thread_num); std::atomic guard_counter(0); std::atomic* guard_counter_ptr = &guard_counter; @@ -1441,15 +1512,16 @@ class AttentionMainLoop { const int64_t q_head_num_stride = input->query_num_heads_stride; const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride; const int64_t kv_cache_block_num_stride = input->cache_num_blocks_stride; - const int32_t sliding_window_left = input->sliding_window_left; - const int32_t sliding_window_right = input->sliding_window_right; + const int32_t sliding_window_size = input->sliding_window_size; const int32_t block_size = input->block_size; const float scale = input->scale; const float softcap_scale = input->softcap; const float* alibi_slopes = input->alibi_slopes; const c10::BFloat16* s_aux = input->s_aux; + const bool* dynamic_causal = input->dynamic_causal; + const bool is_dynamic_causal = dynamic_causal != nullptr; - const bool casual = input->causal; + const bool causal = input->causal; int32_t* const block_table = input->block_table; const int64_t block_table_stride = input->blt_num_tokens_stride; @@ -1532,6 +1604,11 @@ class AttentionMainLoop { &curr_workitem_groups[workitem_group_idx]; const int32_t current_group_idx = current_workitem_group->req_id; + const int32_t current_group_causal = + is_dynamic_causal ? dynamic_causal[current_group_idx] : causal; + auto [sliding_window_left, sliding_window_right] = + AttentionScheduler::calcu_sliding_window_size( + sliding_window_size, current_group_causal); const int32_t kv_start_pos = current_workitem_group->kv_split_pos_start; const int32_t kv_end_pos = current_workitem_group->kv_split_pos_end; @@ -1559,8 +1636,7 @@ class AttentionMainLoop { const int32_t q_end = input->query_start_loc[current_group_idx + 1]; const int32_t q_start = input->query_start_loc[current_group_idx]; const int32_t seq_len = input->seq_lens[current_group_idx]; - const int32_t q_start_pos = - (casual ? seq_len - (q_end - q_start) : 0); + const int32_t q_start_pos = seq_len - (q_end - q_start); const int32_t block_num = (seq_len + block_size - 1) / block_size; // Only apply sink for the first KV split bool use_sink = (s_aux != nullptr && @@ -1610,8 +1686,8 @@ class AttentionMainLoop { const auto [kv_tile_start_pos, kv_tile_end_pos] = AttentionScheduler::calcu_kv_tile_pos( kv_start_pos, kv_end_pos, q_tile_start_pos, - q_tile_end_pos, sliding_window_left, - sliding_window_right); + q_tile_end_pos, sliding_window_size, + current_group_causal); const auto [rounded_kv_tile_start_pos, rounded_kv_tile_end_pos] = AttentionScheduler::align_kv_tile_pos( kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment); @@ -1724,8 +1800,8 @@ class AttentionMainLoop { actual_kv_tile_pos_right] = AttentionScheduler::calcu_kv_tile_pos( kv_tile_pos_left, kv_tile_pos_right, q_tile_pos_left, - q_tile_pos_right, sliding_window_left, - sliding_window_right); + q_tile_pos_right, sliding_window_size, + current_group_causal); const int32_t q_iter_idx = q_head_tile_token_offset / curr_max_q_token_num_per_iter; @@ -1789,7 +1865,7 @@ class AttentionMainLoop { attn_impl.template execute_attention( curr_q_heads_buffer, curr_k_cache, curr_v_cache, logits_buffer, curr_partial_q_buffer, curr_max_buffer, - curr_sum_buffer, curr_block_table, + curr_sum_buffer, curr_block_table, kv_end_pos, aligned_actual_kv_tile_pos_left, aligned_actual_kv_tile_pos_right, actual_kv_token_num, kv_cache_block_num_stride, q_tile_head_num, diff --git a/csrc/cpu/cpu_attn_vsx.hpp b/csrc/cpu/cpu_attn_vsx.hpp index c7e1502bcb0..562a5312571 100644 --- a/csrc/cpu/cpu_attn_vsx.hpp +++ b/csrc/cpu/cpu_attn_vsx.hpp @@ -50,7 +50,16 @@ FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p, b1 = (__vector float)vec_mergel(zeros, raw); } -// Note: c10::Half (FP16) is not supported on PowerPC architecture +// [3] Half (FP16) Specialization +template <> +FORCE_INLINE void load_row8_B_as_f32(const c10::Half* p, + __vector float& b0, + __vector float& b1) { + vec_op::FP16Vec8 fp16_vec(p); + vec_op::FP32Vec8 fp32_vec(fp16_vec); + b0 = fp32_vec.reg.val[0]; + b1 = fp32_vec.reg.val[1]; +} template FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4( @@ -314,8 +323,6 @@ class AttentionImpl { const int64_t num_blocks_stride, const int64_t cache_head_num_stride, const int64_t block_size, const int64_t block_size_stride, const float k_inv = 0.0f, const float v_inv = 0.0f) { - // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on - // PowerPC. The parameters are present to match the common interface. #pragma omp parallel for collapse(2) for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 5839d6c2aaf..07b0aaf8688 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -14,6 +14,18 @@ #define AMX_DISPATCH(...) case cpu_utils::ISA::AMX: #endif +#if defined(ARM_BF16_SUPPORT) + #include "cpu/micro_gemm/cpu_micro_gemm_neon.hpp" + #define NEON_DISPATCH(...) \ + case cpu_utils::ISA::NEON: { \ + using gemm_t = \ + cpu_micro_gemm::MicroGemm; \ + return __VA_ARGS__(); \ + } +#else + #define NEON_DISPATCH(...) case cpu_utils::ISA::NEON: +#endif + #define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \ [&] { \ switch (ISA_TYPE) { \ @@ -23,6 +35,7 @@ cpu_micro_gemm::MicroGemm; \ return __VA_ARGS__(); \ } \ + NEON_DISPATCH(__VA_ARGS__) \ default: { \ TORCH_CHECK(false, "Invalid CPU ISA type."); \ } \ @@ -57,10 +70,12 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, const int32_t input_stride, const int32_t output_stride) { using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; +#if !defined(__aarch64__) // For GPT-OSS interleaved gate-up weights alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30}; vec_op::INT32Vec16 index_vec(index); +#endif vec_op::FP32Vec16 gate_up_max_vec(7.0); vec_op::FP32Vec16 up_min_vec(-7.0); vec_op::FP32Vec16 alpha_vec(1.702); @@ -70,8 +85,15 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, for (int32_t m = 0; m < m_size; ++m) { for (int32_t n = 0; n < n_size; n += 32) { + // Note: AdvSIMD does not support gather loads +#if defined(__aarch64__) + vec_op::FP32Vec16 gate_vec(vec_op::uninit); + vec_op::FP32Vec16 up_vec(vec_op::uninit); + vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); +#else vec_op::FP32Vec16 gate_vec(input + n, index_vec); vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); +#endif gate_vec = gate_vec.min(gate_up_max_vec); up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec); auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec)); @@ -163,7 +185,6 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, vec_op::FP32Vec16 w1_vec(0.7978845608028654); vec_op::FP32Vec16 w2_vec(0.5); vec_op::FP32Vec16 w3_vec(0.044715); - alignas(64) float temp[16]; for (int32_t m = 0; m < m_size; ++m) { for (int32_t n = 0; n < dim; n += 16) { @@ -171,12 +192,9 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, vec_op::FP32Vec16 up_vec(up + n); auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); - - inner_vec.save(temp); - for (int32_t i = 0; i < 16; ++i) { - temp[i] = std::tanh(temp[i]); - } - vec_op::FP32Vec16 tanh_vec(temp); + // Note: can't use fast_exp form because diffusiongemma will generate + // wrong results + auto tanh_vec = inner_vec.tanh(); auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); auto gated_output_fp32 = up_vec * gelu_tanh; scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); @@ -242,13 +260,14 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, constexpr int32_t gemm_n_tile_size = gemm_t::NSize; constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size; + constexpr bool pack_a = gemm_t::PackA; static_assert(gemm_n_tile_size % 16 == 0); TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0); TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); - const int32_t thread_num = omp_get_max_threads(); + const int32_t thread_num = cpu_utils::get_max_threads(); const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( gemm_m_tile_size * input_size_13 * sizeof(scalar_t)); @@ -268,12 +287,18 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, const int32_t w2_input_tile_size = cpu_utils::round_up<64>( gemm_m_tile_size * input_size_2 * sizeof(scalar_t)); + // use w2 input buffer only when we need to pack input + const int32_t w2_input_buffer_size = + pack_a ? cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * + sizeof(scalar_t)) + : 0; const int32_t w2_n_tile_size = [&]() { const int64_t cache_size = cpu_utils::get_available_l2_size(); - // input tile + weight + // input tile + optional packed input + weight const int32_t n_size_cache_limit = - (cache_size - w2_input_tile_size) / (input_size_2 * sizeof(scalar_t)); + (cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) / + (input_size_2 * sizeof(scalar_t)); const int32_t n_size_thread_limit = output_size_2 / std::max(1, thread_num / topk_num); const int32_t n_size = cpu_utils::round_down( @@ -326,6 +351,9 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, const int32_t w13_output_buffer_offset = w13_thread_buffer_offset; w13_thread_buffer_offset += w13_output_buffer_size; + const int32_t w2_input_buffer_offset = w13_thread_buffer_offset; + w13_thread_buffer_offset += w2_input_buffer_size; + // Weighted sum thread buffer const int32_t ws_output_buffer_size = cpu_utils::round_up<64>(output_size_2 * sizeof(float)); @@ -405,7 +433,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, gemm_t gemm; const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t); - const int32_t w13_n_group_stride = 16 * input_size_13; + const int32_t w13_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_13; const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13; for (;;) { @@ -468,8 +497,23 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, token_idx += gemm_m_tile_size) { const int32_t actual_token_num = std::min(gemm_m_tile_size, curr_token_num - token_idx); - // copy inputs - { + + scalar_t* __restrict__ curr_w13_gemm_input_buffer = nullptr; + if constexpr (pack_a) { + // copy and pack inputs + curr_w13_gemm_input_buffer = w13_input_buffer; + const scalar_t* w13_input_rows[gemm_m_tile_size]; + for (int32_t i = 0; i < actual_token_num; ++i) { + w13_input_rows[i] = + input + curr_expand_token_id_buffer[i] * input_size_13; + } + gemm_t::pack_input_from_rows(w13_input_rows, + curr_w13_gemm_input_buffer, + actual_token_num, input_size_13); + curr_expand_token_id_buffer += actual_token_num; + } else { + // copy inputs + curr_w13_gemm_input_buffer = curr_w13_input_buffer; scalar_t* __restrict__ curr_w13_input_buffer_iter = curr_w13_input_buffer; for (int32_t i = 0; i < actual_token_num; ++i) { @@ -501,14 +545,12 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1; scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0; scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1; - scalar_t* __restrict__ curr_w13_input_buffer_iter = - curr_w13_input_buffer; float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer; float* __restrict__ w13_output_buffer_1_iter = w13_output_buffer + actual_n_tile_size / 2; for (int32_t i = 0; i < actual_n_tile_size; i += min_w13_n_tile_size) { - gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_0_iter, + gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_0_iter, w13_output_buffer_0_iter, actual_token_num, input_size_13, input_size_13, w13_n_group_stride, actual_n_tile_size, false); @@ -521,7 +563,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, w13_bias_ptr_0_iter += gemm_n_tile_size; } - gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_1_iter, + gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_1_iter, w13_output_buffer_1_iter, actual_token_num, input_size_13, input_size_13, w13_n_group_stride, actual_n_tile_size, false); @@ -574,7 +616,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, gemm_t gemm; const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2; - const int32_t w2_n_group_stride = 16 * input_size_2; + const int32_t w2_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_2; for (;;) { int32_t task_id = counter_ptr->acquire_counter(); @@ -613,13 +656,30 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, token_idx += gemm_m_tile_size) { const int32_t actual_token_num = std::min(gemm_m_tile_size, curr_token_num - token_idx); + scalar_t* __restrict__ curr_w2_gemm_input_buffer = + curr_w13_gemm_output_buffer; + if constexpr (pack_a) { + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * w13_thread_buffer_offset; + scalar_t* __restrict__ w2_input_buffer = + reinterpret_cast(thread_buffer + + w2_input_buffer_offset); + curr_w2_gemm_input_buffer = w2_input_buffer; + const scalar_t* w2_input_rows[gemm_m_tile_size]; + for (int32_t i = 0; i < actual_token_num; ++i) { + w2_input_rows[i] = curr_w13_gemm_output_buffer + i * input_size_2; + } + gemm_t::pack_input_from_rows(w2_input_rows, + curr_w2_gemm_input_buffer, + actual_token_num, input_size_2); + } scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr; scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr; float* __restrict__ curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer; for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) { - gemm.gemm(curr_w13_gemm_output_buffer, w2_weight_ptr_iter, + gemm.gemm(curr_w2_gemm_input_buffer, w2_weight_ptr_iter, curr_w2_gemm_output_buffer_iter, actual_token_num, input_size_2, input_size_2, w2_n_group_stride, output_size_2, false); diff --git a/csrc/cpu/cpu_tanhf_neon.hpp b/csrc/cpu/cpu_tanhf_neon.hpp new file mode 100644 index 00000000000..2ea7f336513 --- /dev/null +++ b/csrc/cpu/cpu_tanhf_neon.hpp @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_TANHF_NEON_HPP +#define CPU_TANHF_NEON_HPP + +#include +#include + +namespace vec_op { + +namespace { + +struct TanhfConstants { + float32x4_t special_bound; + float32x4_t two; + float32x4_t c0; + float32x4_t c2; + int32x4_t exponent_bias; + float c1; + float c3; + float two_over_ln2; + float c4; + float ln2_hi; + float ln2_lo; +}; + +const TanhfConstants kTanhfConstants = { + // 9.01, above which tanhf rounds to 1 (or -1 for negative). + .special_bound = vdupq_n_f32(0x1.205966p+3f), + .two = vdupq_n_f32(0x1.0p+1f), + .c0 = vdupq_n_f32(0x1.fffffep-2f), + .c2 = vdupq_n_f32(0x1.555736p-5f), + .exponent_bias = vdupq_n_s32(0x3f800000), + .c1 = 0x1.5554aep-3f, + .c3 = 0x1.12287cp-7f, + .two_over_ln2 = 0x1.715476p+1f, + .c4 = 0x1.6b55a2p-10f, + .ln2_hi = 0x1.62e4p-1f, + .ln2_lo = 0x1.7f7d1cp-20f, +}; + +// Return the ptr but hide it's value from the compiler so accesses +// through it can't be optimised based on contents. +template +inline const T* ptr_barrier(const T* ptr) { + const T* opaque_ptr = ptr; + __asm__("" : "+r"(opaque_ptr)); + return opaque_ptr; +} + +// Check whether any lanes in the mask are set +inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; } + +// e^2x - 1 inline helper +inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) { + float32x2_t ln2 = vld1_f32(&d->ln2_hi); + float32x4_t lane_consts = vld1q_f32(&d->c1); + + // Reduce argument: f in [-ln2/2, ln2/2], i is exact. + float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2)); + int32x4_t i = vcvtq_s32_f32(j); + float32x4_t f = vaddq_f32(x, x); + f = vfmsq_lane_f32(f, j, ln2, 0); + f = vfmsq_lane_f32(f, j, ln2, 1); + + // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f) + float32x4_t f2 = vmulq_f32(f, f); + float32x4_t f4 = vmulq_f32(f2, f2); + float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0); + float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1); + float32x4_t poly = vfmaq_f32(p01, f2, p23); + poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3); + poly = vfmaq_f32(f, f2, poly); + + // scale = 2^i + int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias); + float32x4_t scale = vreinterpretq_f32_s32(u); + return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale); +} + +// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1 +inline float32x4_t special_case(float32x4_t x, float32x4_t q, + uint32x4_t special) { + const TanhfConstants* d = ptr_barrier(&kTanhfConstants); + + float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two)); + uint32x4_t ix = vreinterpretq_u32_f32(x); + uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias); + uint32x4_t sign_mask = vdupq_n_u32(0x80000000u); + uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits); + float32x4_t special_y = vreinterpretq_f32_u32(special_bits); + return vbslq_f32(special, special_y, y); +} + +} // namespace + +// Implementation of tanhf adapted from Arm Optimized Routines (tanhf +// AdvSIMD) +// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c +// +// Approximation for single-precision vector tanh(x), using a simplified +// version of expm1f. The maximum error is 2.08 + 0.5 ULP: +// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5. +inline float32x4_t fast_tanhf_f32x4(float32x4_t x) { + const TanhfConstants* d = ptr_barrier(&kTanhfConstants); + + // tanh(x) = (e^2x - 1) / (e^2x + 1) + // q = e^2x -1 + float32x4_t q = e2xm1f_inline(x, d); + + // Check for special cases + uint32x4_t special = vcagtq_f32(x, d->special_bound); + + // Fall back to vectorised special case for any lanes which would cause + // expm1 to overflow + if (any_u32(special)) { + return special_case(x, q, special); + } + + // Complete fast path if no special lanes + // tanh(x) = q / (q+2) + return vdivq_f32(q, vaddq_f32(q, d->two)); +} + +} // namespace vec_op + +#endif // CPU_TANHF_NEON_HPP \ No newline at end of file diff --git a/csrc/cpu/cpu_types.hpp b/csrc/cpu/cpu_types.hpp index 744c80c8f53..50ef5e10b83 100644 --- a/csrc/cpu/cpu_types.hpp +++ b/csrc/cpu/cpu_types.hpp @@ -4,7 +4,7 @@ #if defined(__x86_64__) // x86 implementation #include "cpu_types_x86.hpp" -#elif defined(__POWER9_VECTOR__) +#elif defined(__powerpc__) // ppc implementation #include "cpu_types_vsx.hpp" #elif defined(__s390x__) @@ -25,4 +25,20 @@ #include #endif -#endif \ No newline at end of file +#include + +namespace cpu_utils { +// Without OpenMP the omp pragmas compile to serial loops, so report 1: kernels +// that barrier on the thread count would otherwise deadlock. +inline int get_max_threads() { +#ifdef _OPENMP + return omp_get_max_threads(); +#else + TORCH_WARN_ONCE( + "vLLM CPU was built without OpenMP; running single-threaded."); + return 1; +#endif +} +} // namespace cpu_utils + +#endif diff --git a/csrc/cpu/cpu_types_arm.hpp b/csrc/cpu/cpu_types_arm.hpp index b408731f40d..294dee90bd8 100644 --- a/csrc/cpu/cpu_types_arm.hpp +++ b/csrc/cpu/cpu_types_arm.hpp @@ -3,6 +3,8 @@ #include +#include "cpu/cpu_tanhf_neon.hpp" + #include #include #include @@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper { explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {}; explicit FP32Vec4(const FP32Vec4& data) : Base(data) {}; + + FORCE_INLINE FP32Vec4 tanh() const { + return FP32Vec4(fast_tanhf_f32x4(reg.val[0])); + } }; struct FP32Vec8 : public VectorizedRegWrapper { @@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper { reg.val[1] = Vectorized(data.val[1]); } + FORCE_INLINE FP32Vec8 tanh() const { + FP32Vec8 r(uninit); + r.reg.val[0] = Vectorized(fast_tanhf_f32x4(reg.val[0])); + r.reg.val[1] = Vectorized(fast_tanhf_f32x4(reg.val[1])); + return r; + } + FORCE_INLINE float reduce_sum() const noexcept { float answer = 0; std::plus add; @@ -497,6 +510,35 @@ struct FP32Vec16 : public VectorizedRegWrapper { reg.val[3] = Vectorized(vcvt_f32_f16(vget_high_f16(v.reg.val[1]))); }; + FORCE_INLINE FP32Vec16 tanh() const { + FP32Vec16 r(uninit); + r.reg.val[0] = Vectorized(fast_tanhf_f32x4(reg.val[0])); + r.reg.val[1] = Vectorized(fast_tanhf_f32x4(reg.val[1])); + r.reg.val[2] = Vectorized(fast_tanhf_f32x4(reg.val[2])); + r.reg.val[3] = Vectorized(fast_tanhf_f32x4(reg.val[3])); + return r; + } + + static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even, + FP32Vec16& odd) noexcept { + const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4)); + const float32x4x2_t x23 = + vuzpq_f32(vld1q_f32(ptr + 8), vld1q_f32(ptr + 12)); + const float32x4x2_t x45 = + vuzpq_f32(vld1q_f32(ptr + 16), vld1q_f32(ptr + 20)); + const float32x4x2_t x67 = + vuzpq_f32(vld1q_f32(ptr + 24), vld1q_f32(ptr + 28)); + + even.reg.val[0] = VectorizedT(x01.val[0]); + even.reg.val[1] = VectorizedT(x23.val[0]); + even.reg.val[2] = VectorizedT(x45.val[0]); + even.reg.val[3] = VectorizedT(x67.val[0]); + odd.reg.val[0] = VectorizedT(x01.val[1]); + odd.reg.val[1] = VectorizedT(x23.val[1]); + odd.reg.val[2] = VectorizedT(x45.val[1]); + odd.reg.val[3] = VectorizedT(x67.val[1]); + } + FORCE_INLINE FP32Vec16 operator+(const FP32Vec16& b) const noexcept { FP32Vec16 r(uninit); r.reg.val[0] = reg.val[0] + b.reg.val[0]; @@ -515,6 +557,15 @@ struct FP32Vec16 : public VectorizedRegWrapper { return r; } + FORCE_INLINE FP32Vec16 operator-() const noexcept { + FP32Vec16 r(uninit); + r.reg.val[0] = reg.val[0].neg(); + r.reg.val[1] = reg.val[1].neg(); + r.reg.val[2] = reg.val[2].neg(); + r.reg.val[3] = reg.val[3].neg(); + return r; + } + FORCE_INLINE FP32Vec16 operator*(const FP32Vec16& b) const noexcept { FP32Vec16 r(uninit); r.reg.val[0] = reg.val[0] * b.reg.val[0]; @@ -933,4 +984,4 @@ inline void storeFP32(float v, c10::BFloat16* ptr) { inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }; -}; // namespace vec_op \ No newline at end of file +}; // namespace vec_op diff --git a/csrc/cpu/cpu_types_riscv_defs.hpp b/csrc/cpu/cpu_types_riscv_defs.hpp index 8871617f05f..16475505d9f 100644 --- a/csrc/cpu/cpu_types_riscv_defs.hpp +++ b/csrc/cpu/cpu_types_riscv_defs.hpp @@ -3,13 +3,17 @@ // 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 +// LMUL_ expands to the LMUL suffix giving N total bits of vector data. +// LMUL_64 is used by 8-lane int8/uint8 vectors. +// VLEN=128: +// LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8 +// VLEN=256: +// LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4 #include #if __riscv_v_min_vlen == 128 + #define LMUL_64 mf2 #define LMUL_128 m1 #define LMUL_256 m2 #define LMUL_512 m4 @@ -17,6 +21,7 @@ #define BOOL_256 b16 #define BOOL_512 b8 #elif __riscv_v_min_vlen == 256 + #define LMUL_64 mf4 #define LMUL_128 mf2 #define LMUL_256 m1 #define LMUL_512 m2 @@ -41,6 +46,16 @@ // ---- Semantic fixed-vector typedefs (named by element count) ---- +// uint8 / int8 +typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t + __attribute__((riscv_rvv_vector_bits(64))); +typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t + __attribute__((riscv_rvv_vector_bits(64))); + +// int16 +typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t + __attribute__((riscv_rvv_vector_bits(128))); + // float16 typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t __attribute__((riscv_rvv_vector_bits(128))); @@ -57,6 +72,10 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t __attribute__((riscv_rvv_vector_bits(1024))); +// int8 +typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t + __attribute__((riscv_rvv_vector_bits(128))); + // int32 typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t __attribute__((riscv_rvv_vector_bits(256))); diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index 06a38c780a2..70cb0ab52de 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -9,10 +9,14 @@ #include #include +#include #include #include #include #include + +#include "float_convert.hpp" + namespace vec_op { // FP8 KV cache is not supported on RISC-V. These tag types and the @@ -210,11 +214,18 @@ struct BF16Vec32 : public Vec { 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); + // Widen LMUL_128 → LMUL_256 so vslideup operands share a type. + // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2. + fixed_u16x16_t ext = + RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val); + // Build 16-element half: place the 8 elements at offsets 0 and 8. + fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16); + // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128). + fixed_u16x32_t dst = + RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst); }; void save(void* ptr) const { @@ -245,8 +256,7 @@ struct BF16Vec8 : public Vec { 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); + tmp[i] = bf16_to_float(u16[i]); } reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8); } @@ -256,9 +266,7 @@ struct BF16Vec8 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } void save(void* ptr, int elem_num) const { @@ -266,9 +274,7 @@ struct BF16Vec8 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } void save_strided(void* ptr, ptrdiff_t stride) const { @@ -277,10 +283,8 @@ struct BF16Vec8 : public Vec { 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; + *reinterpret_cast(u8 + i * byte_stride) = + float_to_bf16(tmp[i]); } } }; @@ -292,8 +296,7 @@ struct BF16Vec16 : public Vec { 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); + tmp[i] = bf16_to_float(u16[i]); } reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16); } @@ -306,9 +309,7 @@ struct BF16Vec16 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } void save(void* ptr, int elem_num) const { @@ -316,9 +317,7 @@ struct BF16Vec16 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } void save_strided(void* ptr, ptrdiff_t stride) const { @@ -327,10 +326,8 @@ struct BF16Vec16 : public Vec { 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; + *reinterpret_cast(u8 + i * byte_stride) = + float_to_bf16(tmp[i]); } } }; @@ -343,8 +340,7 @@ struct BF16Vec32 : public Vec { 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); + tmp[i] = bf16_to_float(u16[i]); } reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32); } @@ -371,9 +367,7 @@ struct BF16Vec32 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } @@ -382,9 +376,7 @@ struct BF16Vec32 : public Vec { 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); + u16[i] = float_to_bf16(tmp[i]); } } @@ -394,10 +386,8 @@ struct BF16Vec32 : public Vec { 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; + *reinterpret_cast(u8 + i * byte_stride) = + float_to_bf16(tmp[i]); } } }; @@ -640,17 +630,29 @@ struct FP32Vec16 : public Vec { data.reg, data.reg)) {}; explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { - const uint64_t q_values = static_cast(value); - auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM); - auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM); - auto shifts = - RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM); - auto shifted = - RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM); - auto idx64 = - RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM); - auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM); - reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM); + // Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on + // VLEN=128 / m4 on VLEN=256), which causes heavy register spilling. + constexpr int HALF = VEC_ELEM_NUM / 2; + const auto q = static_cast(value); + const uint32_t lo = static_cast(q); + const uint32_t hi = static_cast(q >> 32); + + auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF); + auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF); + + auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF); + auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF, + HALF); + + auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF); + auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF, + HALF); + + auto idx = + RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi); + reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM); } explicit FP32Vec16(const FP16Vec16& v); @@ -734,10 +736,18 @@ struct FP32Vec16 : public Vec { return FP32Vec16( RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); } + FP32Vec16 max(const FP32Vec16& b, const int elem_num) const { + return FP32Vec16( + RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num)); + } FP32Vec16 min(const FP32Vec16& b) const { return FP32Vec16( RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); } + FP32Vec16 min(const FP32Vec16& b, const int elem_num) const { + return FP32Vec16( + RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num)); + } FP32Vec16 abs() const { return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM)); } @@ -867,6 +877,27 @@ struct FP32Vec16 : public Vec { } }; +struct INT8Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + fixed_i8x16_t reg; + + explicit INT8Vec16(const FP32Vec16& vec) { + auto i32_vec = + RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM); + auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)( + i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM); + reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU, + VEC_ELEM_NUM); + } + + void save(int8_t* ptr) const { + RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM); + } + void save(int8_t* ptr, int elem_num) const { + RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num); + } +}; + // ============================================================================ // Type Traits & Global Helpers // ============================================================================ @@ -956,9 +987,7 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) #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); + *reinterpret_cast(ptr) = float_to_bf16(v); } inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {} inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {} diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp index d1c2fc85933..94b5179b171 100644 --- a/csrc/cpu/cpu_types_scalar.hpp +++ b/csrc/cpu/cpu_types_scalar.hpp @@ -363,6 +363,13 @@ struct FP32Vec16 : public Vec { return FP32Vec16(ret); } + FP32Vec16 tanh() const { + f32x16_t ret; + unroll_loop( + [&ret, this](int i) { ret.val[i] = std::tanh(reg.val[i]); }); + return FP32Vec16(ret); + } + float reduce_sum() const { float result = 0.0f; unroll_loop( diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index ba65e27a15e..250c870dbe4 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -13,10 +13,10 @@ namespace vec_op { struct fp8_e4m3_tag {}; struct fp8_e5m2_tag {}; -// FIXME: FP16 is not fully supported in Torch-CPU -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) @@ -34,6 +34,87 @@ struct fp8_e5m2_tag {}; #define FORCE_INLINE __attribute__((always_inline)) inline namespace { + +FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) { + const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000}; + const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00}; + const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF}; + const __vector unsigned int bias_adj = {112, 112, 112, 112}; + const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F}; + const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF}; + + __vector unsigned int s = (x & mask_sign) << 16; + __vector unsigned int e = (x & mask_exp) >> 10; + __vector unsigned int m = (x & mask_mant) << 13; + + __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16); + + __vector unsigned int e_normal = e + bias_adj; + e = vec_sel(e_normal, exp_max_fp32, is_nan_inf); + + return (__vector float)(s | (e << 23) | m); +} + +FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) { + __vector unsigned int in = (__vector unsigned int)f_in; + + const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000, + 0x80000000, 0x80000000}; + const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000, + 0x7F800000}; + const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF, + 0x007FFFFF, 0x007FFFFF}; + + const __vector signed int bias_adj = {112, 112, 112, 112}; + const __vector signed int zero = {0, 0, 0, 0}; + const __vector signed int max_exp = {31, 31, 31, 31}; + const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF}; + const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F}; + + __vector unsigned int s = (in & mask_sign_32) >> 16; + __vector unsigned int e_u = (in & mask_exp_32) >> 23; + + __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32); + + __vector signed int e_s = (__vector signed int)e_u; + e_s = vec_sub(e_s, bias_adj); + e_s = vec_max(e_s, zero); + e_s = vec_min(e_s, max_exp); + __vector unsigned int e_normal = (__vector unsigned int)e_s; + + __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf); + + const __vector unsigned int one_v = {1, 1, 1, 1}; + const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF}; + + __vector unsigned int round_bit = (in >> 12) & one_v; + __vector unsigned int sticky = in & mask_sticky; + __vector unsigned int m = (in & mask_mant_32) >> 13; + __vector unsigned int lsb = m & one_v; + + // Round up if: round_bit && (sticky || lsb) + __vector __bool int sticky_nonzero = + vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0}); + __vector __bool int lsb_set = vec_cmpeq(lsb, one_v); + __vector __bool int round_up = + vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set)); + + m = vec_sel(m, m + one_v, round_up); + + const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF}; + const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E}; + __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask); + __vector __bool int would_overflow_to_inf = + vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp)); + __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16); + e_final = vec_sel(e_final, e_inc, mant_overflows); + m = vec_and(m, mant_mask); + e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf); + m = vec_sel(m, mant_mask, would_overflow_to_inf); + + return s | (e_final << 10) | m; +} + template constexpr void unroll_loop_item(std::integer_sequence, F&& f) { (f(std::integral_constant{}), ...); @@ -89,6 +170,19 @@ struct BF16Vec8 : public Vec { } }; +struct FP16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + + __vector signed short reg; + + explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {} + explicit FP16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + *reinterpret_cast<__vector signed short*>(ptr) = reg; + } +}; + struct FP16Vec16 : public Vec { constexpr static int VEC_ELEM_NUM = 16; ss16x8x2_t reg; @@ -124,13 +218,11 @@ struct BF16Vec16 : public Vec { ss16x8x2_t reg; explicit BF16Vec16(const void* ptr) { - // Load 256 bits in two parts reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr); reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr); } explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {} - explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { @@ -142,20 +234,16 @@ struct BF16Vec16 : public Vec { void save(void* ptr, const int elem_num) const { const int clamped_elem = std::max(0, std::min(elem_num, 16)); - // Calculate elements to store in each 128-bit part (8 elements each) const int elements_val0 = std::min(clamped_elem, 8); const int elements_val1 = std::max(clamped_elem - 8, 0); - // Convert elements to bytes (2 bytes per element) const size_t bytes_val0 = elements_val0 * sizeof(signed short); const size_t bytes_val1 = elements_val1 * sizeof(signed short); signed short* dest = static_cast(ptr); - // Store the first part using vec_xst_len if (bytes_val0 > 0) { vec_xst_len(reg.val[0], dest, bytes_val0); } - // Store the second part if needed if (bytes_val1 > 0) { vec_xst_len(reg.val[1], dest + elements_val0, bytes_val1); } @@ -238,6 +326,15 @@ struct FP32Vec8 : public Vec { reg.val[1] = (__vector float)vec_mergel(zero, v.reg); } + explicit FP32Vec8(const FP16Vec8& v) { + __vector unsigned short raw_u = (__vector unsigned short)v.reg; + __vector unsigned int raw_hi = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u); + __vector unsigned int raw_lo = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u); + reg.val[0] = fp16_to_fp32_bits(raw_hi); + reg.val[1] = fp16_to_fp32_bits(raw_lo); + } float reduce_sum() const { AliasReg ar; ar.reg = reg; @@ -247,53 +344,133 @@ struct FP32Vec8 : public Vec { return result; } - FP32Vec8 exp() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::exp(ar.values[0]); - ret.val[0][1] = std::exp(ar.values[1]); - ret.val[0][2] = std::exp(ar.values[2]); - ret.val[0][3] = std::exp(ar.values[3]); - ret.val[1][0] = std::exp(ar.values[4]); - ret.val[1][1] = std::exp(ar.values[5]); - ret.val[1][2] = std::exp(ar.values[6]); - ret.val[1][3] = std::exp(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + f32x4x2_t out; + const __vector float log2e = vec_splats(1.44269504088896341f); + const __vector float one = vec_splats(1.0f); + const __vector float min_x = vec_splats(-87.3f); + const __vector float max_x = vec_splats(88.7f); + + // 5th-degree minimax polynomial for 2^r (r in [0,1)) + const __vector float c1 = vec_splats(0.6931471805599453f); + const __vector float c2 = vec_splats(0.240226506959101f); + const __vector float c3 = vec_splats(0.05550410866482158f); + const __vector float c4 = vec_splats(0.009618129107628477f); + const __vector float c5 = vec_splats(0.0013333558146428443f); + + for (int i = 0; i < 2; i++) { + __vector float x = reg.val[i]; + x = vec_max(x, min_x); + x = vec_min(x, max_x); + + __vector float y = vec_mul(x, log2e); + + __vector float kf = vec_floor(y); + __vector float r = vec_sub(y, kf); + + // Convert float to signed integer. Use vec_cts for PowerPC AltiVec + // compatibility. + __vector signed int k = vec_cts(kf, 0); + const __vector signed int min_k = vec_splats((signed int)-126); + const __vector signed int max_k = vec_splats((signed int)127); + k = vec_min(vec_max(k, min_k), max_k); + + // Build 2^k from exponent bits + __vector signed int exp_int = vec_add(k, vec_splats((signed int)127)); + __vector unsigned int bits = (__vector unsigned int)exp_int; + bits = vec_sl(bits, vec_splats((unsigned int)23)); + __vector float pow2k = (__vector float)bits; + + // Improved minimax polynomial + __vector float poly = vec_madd(c5, r, c4); + poly = vec_madd(poly, r, c3); + poly = vec_madd(poly, r, c2); + poly = vec_madd(poly, r, c1); + poly = vec_madd(poly, r, one); + + out.val[i] = vec_mul(pow2k, poly); + } + return FP32Vec8(out); } FP32Vec8 tanh() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::tanh(ar.values[0]); - ret.val[0][1] = std::tanh(ar.values[1]); - ret.val[0][2] = std::tanh(ar.values[2]); - ret.val[0][3] = std::tanh(ar.values[3]); - ret.val[1][0] = std::tanh(ar.values[4]); - ret.val[1][1] = std::tanh(ar.values[5]); - ret.val[1][2] = std::tanh(ar.values[6]); - ret.val[1][3] = std::tanh(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + const __vector float one = vec_splats(1.0f); + const __vector float two = vec_splats(2.0f); + const __vector float zero = vec_splats(0.0f); + const __vector float sat = vec_splats(9.0f); + + f32x4x2_t out; + + for (int i = 0; i < 2; i++) { + __vector float x = reg.val[i]; + __vector float ax = vec_abs(x); + + __vector bool int mask = vec_cmpge(x, zero); + __vector float sign = vec_sel(vec_splats(-1.0f), one, mask); + + __vector bool int saturated = vec_cmpge(ax, sat); + + __vector float two_x = vec_mul(x, two); + f32x4x2_t tmp; + tmp.val[0] = two_x; + tmp.val[1] = two_x; + FP32Vec8 temp_vec(tmp); + vector float e = temp_vec.exp().reg.val[0]; + + vector float num = vec_sub(e, one); + vector float den = vec_add(e, one); + vector float t = vec_div(num, den); + + out.val[i] = vec_sel(t, sign, saturated); + } + return FP32Vec8(out); } FP32Vec8 er() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::erf(ar.values[0]); - ret.val[0][1] = std::erf(ar.values[1]); - ret.val[0][2] = std::erf(ar.values[2]); - ret.val[0][3] = std::erf(ar.values[3]); - ret.val[1][0] = std::erf(ar.values[4]); - ret.val[1][1] = std::erf(ar.values[5]); - ret.val[1][2] = std::erf(ar.values[6]); - ret.val[1][3] = std::erf(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + const vector float a1 = vec_splats(0.254829592f); + const vector float a2 = vec_splats(-0.284496736f); + const vector float a3 = vec_splats(1.421413741f); + const vector float a4 = vec_splats(-1.453152027f); + const vector float a5 = vec_splats(1.061405429f); + const vector float p = vec_splats(0.3275911f); + const vector float one = vec_splats(1.0f); + const vector float zero = vec_splats(0.0f); + const vector float sat = vec_splats(6.0f); + + f32x4x2_t ret; + + for (int i = 0; i < 2; i++) { + vector float x = reg.val[i]; + vector float ax = vec_abs(x); + + vector bool int mask = vec_cmpge(x, zero); + vector float sign = vec_sel(vec_splats(-1.0f), one, mask); + + vector bool int saturated = vec_cmpge(ax, sat); + + vector float t = vec_div(one, vec_madd(p, ax, one)); + + vector float poly = a5; + poly = vec_madd(poly, t, a4); + poly = vec_madd(poly, t, a3); + poly = vec_madd(poly, t, a2); + poly = vec_madd(poly, t, a1); + poly = vec_mul(poly, t); + + vector float x_squared = vec_mul(x, x); + vector float neg_x_squared = vec_mul(vec_splats(-1.0f), x_squared); + f32x4x2_t tmp; + tmp.val[0] = neg_x_squared; + tmp.val[1] = neg_x_squared; + FP32Vec8 exp_input(tmp); + vector float exp_term = exp_input.exp().reg.val[0]; + + vector float y = vec_nmsub(poly, exp_term, one); + vector float erf_val = vec_mul(sign, y); + + ret.val[i] = vec_sel(erf_val, sign, saturated); + } + return FP32Vec8(ret); } FP32Vec8 operator*(const FP32Vec8& b) const { @@ -410,8 +587,9 @@ struct FP32Vec16 : public Vec { reg.val[3] = vec_xl(48, ptr); } + explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {} + explicit FP32Vec16(const FP16Vec16&); explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {} - explicit FP32Vec16(f32x4x4_t data) : reg(data) {} explicit FP32Vec16(const FP32Vec16& data) { @@ -435,7 +613,6 @@ struct FP32Vec16 : public Vec { reg.val[3] = data.reg.val[1]; } - explicit FP32Vec16(const FP16Vec16& v); explicit FP32Vec16(const BF16Vec16& v) { reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]); reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]); @@ -502,28 +679,20 @@ struct FP32Vec16 : public Vec { FP32Vec16 max(const FP32Vec16& b, int elem_num) const { FP32Vec16 result; - // Create a vector of element indices for each chunk __vector unsigned int indices = {0, 1, 2, 3}; __vector unsigned int elem_num_vec = vec_splats(static_cast(elem_num)); - // Compute masks for each chunk - __vector unsigned int chunk_offset0 = {0, 0, 0, - 0}; // Chunk 0: Elements 0-3 - __vector unsigned int chunk_offset1 = {4, 4, 4, - 4}; // Chunk 1: Elements 4-7 - __vector unsigned int chunk_offset2 = {8, 8, 8, - 8}; // Chunk 2: Elements 8-11 - __vector unsigned int chunk_offset3 = {12, 12, 12, - 12}; // Chunk 3: Elements 12-15 + __vector unsigned int chunk_offset0 = {0, 0, 0, 0}; + __vector unsigned int chunk_offset1 = {4, 4, 4, 4}; + __vector unsigned int chunk_offset2 = {8, 8, 8, 8}; + __vector unsigned int chunk_offset3 = {12, 12, 12, 12}; - // Compute masks for each chunk __vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec); __vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec); __vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec); __vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec); - // Apply masks to compute the result for each chunk result.reg.val[0] = vec_sel(this->reg.val[0], vec_max(this->reg.val[0], b.reg.val[0]), mask0); result.reg.val[1] = vec_sel(this->reg.val[1], @@ -626,6 +795,16 @@ struct FP32Vec16 : public Vec { vec_xst(reg.val[3], 48, ptr); } + void save(c10::Half* ptr) const { + FP16Vec16 fp16_vec(*this); + fp16_vec.save(ptr); + } + + void save(c10::Half* ptr, const int elem_num) const { + FP16Vec16 fp16_vec(*this); + fp16_vec.save(ptr, elem_num); + } + void save(float* ptr, const int elem_num) const { const int elements_in_chunk1 = (elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0; @@ -659,7 +838,7 @@ struct FP32Vec16 : public Vec { }; struct INT8Vec16 : public Vec { - constexpr static int VEC_NUM_ELEM = 16; // 128 bits / 8 bits = 16 + constexpr static int VEC_NUM_ELEM = 16; union AliasReg { __vector signed char reg; @@ -707,6 +886,11 @@ struct VecType { using vec_type = BF16Vec8; }; +template <> +struct VecType { + using vec_type = FP16Vec8; +}; + template void storeFP32(float v, T* ptr) { *ptr = v; @@ -723,6 +907,15 @@ inline void storeFP32(float v, c10::BFloat16* ptr) { *ptr = *(v_ptr + 1); } +template <> +inline void storeFP32(float v, c10::Half* ptr) { + __vector float v_vec = {v, 0.0f, 0.0f, 0.0f}; + __vector unsigned int fp16_bits = fp32_to_fp16_bits(v_vec); + unsigned short result = + (unsigned short)((__vector unsigned short)fp16_bits)[0]; + *reinterpret_cast(ptr) = result; +} + #ifndef __VEC_CLASS_FP_NAN #define __VEC_CLASS_FP_NAN (1 << 6) #endif @@ -769,38 +962,39 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) { #endif } +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { + __vector unsigned int fp16_hi = fp32_to_fp16_bits(v.reg.val[0]); + __vector unsigned int fp16_lo = fp32_to_fp16_bits(v.reg.val[1]); + reg = (__vector signed short)vec_perm((__vector unsigned char)fp16_hi, + (__vector unsigned char)fp16_lo, omask); +} + inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { - alignas(16) float temp_fp32[16]; - alignas(16) c10::Half temp_fp16[16]; - - vec_xst(v.reg.val[0], 0, temp_fp32); - vec_xst(v.reg.val[1], 16, temp_fp32); - vec_xst(v.reg.val[2], 32, temp_fp32); - vec_xst(v.reg.val[3], 48, temp_fp32); - - for (int i = 0; i < 16; i++) { - temp_fp16[i] = c10::Half(temp_fp32[i]); - } - - reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16); - reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16); + __vector unsigned int fp16_0 = fp32_to_fp16_bits(v.reg.val[0]); + __vector unsigned int fp16_1 = fp32_to_fp16_bits(v.reg.val[1]); + __vector unsigned int fp16_2 = fp32_to_fp16_bits(v.reg.val[2]); + __vector unsigned int fp16_3 = fp32_to_fp16_bits(v.reg.val[3]); + reg.val[0] = (__vector signed short)vec_perm( + (__vector unsigned char)fp16_0, (__vector unsigned char)fp16_1, omask); + reg.val[1] = (__vector signed short)vec_perm( + (__vector unsigned char)fp16_2, (__vector unsigned char)fp16_3, omask); } inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { - alignas(16) c10::Half temp_fp16[16]; - alignas(16) float temp_fp32[16]; - - vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16); - vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16); - - for (int i = 0; i < 16; i++) { - temp_fp32[i] = float(temp_fp16[i]); - } - - reg.val[0] = vec_xl(0, temp_fp32); - reg.val[1] = vec_xl(16, temp_fp32); - reg.val[2] = vec_xl(32, temp_fp32); - reg.val[3] = vec_xl(48, temp_fp32); + __vector unsigned short raw_u0 = (__vector unsigned short)v.reg.val[0]; + __vector unsigned short raw_u1 = (__vector unsigned short)v.reg.val[1]; + __vector unsigned int raw_hi0 = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u0); + __vector unsigned int raw_lo0 = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u0); + __vector unsigned int raw_hi1 = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u1); + __vector unsigned int raw_lo1 = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u1); + reg.val[0] = fp16_to_fp32_bits(raw_hi0); + reg.val[1] = fp16_to_fp32_bits(raw_lo0); + reg.val[2] = fp16_to_fp32_bits(raw_hi1); + reg.val[3] = fp16_to_fp32_bits(raw_lo1); } inline BF16Vec16::BF16Vec16(const FP32Vec16& v) { @@ -864,7 +1058,6 @@ inline void prefetch(const void* addr) { struct INT8Vec64 { __vector signed char data[4]; - INT8Vec64() = default; explicit INT8Vec64(const int8_t* ptr) { @@ -900,5 +1093,4 @@ struct INT8Vec64 { void nt_save(int8_t* ptr) const { save(ptr); } }; } // namespace vec_op - #endif diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index 2e0af466b64..bf96554a8df 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -3,7 +3,9 @@ #define CPU_TYPES_VXE_HPP #include +#include #include +#include #include #include namespace vec_op { @@ -817,8 +819,7 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) { // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can // produce incorrect results for some inputs. Process each of the 4 vectors // separately. - uint32_t in; - std::memcpy(&in, &v, sizeof(in)); + uint32_t in = std::bit_cast(v); uint32_t s = (in & 0x80000000) >> 16; // Sign uint32_t e = (in & 0x7F800000) >> 23; // Exponent diff --git a/csrc/cpu/cpu_types_x86.hpp b/csrc/cpu/cpu_types_x86.hpp index 396b9b7e041..d2a72ce9ccd 100644 --- a/csrc/cpu/cpu_types_x86.hpp +++ b/csrc/cpu/cpu_types_x86.hpp @@ -3,6 +3,7 @@ #define CPU_TYPES_X86_HPP #include +#include #include #ifndef __AVX2__ @@ -592,6 +593,8 @@ struct FP32Vec16 : public Vec { FP32Vec16 abs() const { return FP32Vec16(_mm512_abs_ps(reg)); } + FP32Vec16 tanh() const { return FP32Vec16(Sleef_tanhf16_u10(reg)); } + float reduce_sum() const { return _mm512_reduce_add_ps(reg); } float reduce_max() const { return _mm512_reduce_max_ps(reg); } @@ -789,6 +792,12 @@ struct FP32Vec16 : public Vec { _mm256_andnot_ps(sign_mask, reg_high)); } + FP32Vec16 tanh() const { + FP32Vec8 low(reg_low); + FP32Vec8 high(reg_high); + return FP32Vec16(low.tanh().reg, high.tanh().reg); + } + FP32Vec16 min(const FP32Vec16& b) const { return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low), _mm256_min_ps(reg_high, b.reg_high)); diff --git a/csrc/cpu/cpu_wna16.cpp b/csrc/cpu/cpu_wna16.cpp index 533f2096354..ae7aef74c44 100644 --- a/csrc/cpu/cpu_wna16.cpp +++ b/csrc/cpu/cpu_wna16.cpp @@ -4,6 +4,9 @@ #ifdef CPU_CAPABILITY_AMXBF16 #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp" #endif +#if defined(__riscv_v) + #include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp" +#endif #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp" #define VLLM_DISPATCH_CASE_16B_TYPES(...) \ @@ -152,7 +155,7 @@ void cpu_gemm_wna16_impl( constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; constexpr int32_t n_block_size = 16; static_assert(gemm_n_tile_size % n_block_size == 0); - const int32_t thread_num = omp_get_max_threads(); + const int32_t thread_num = cpu_utils::get_max_threads(); // a simple schedule policy, just to hold more B tiles in L2 and make sure // each thread has tasks @@ -319,6 +322,8 @@ void cpu_gemm_wna16( return ISA::AMX; } else if (isa_hint == "vec") { return ISA::VEC; + } else if (isa_hint == "rvv") { + return ISA::RVV; } else { TORCH_CHECK(false, "unsupported isa hint: " + isa_hint); } @@ -397,6 +402,40 @@ void cpu_gemm_wna16( pack_factor); return; } + } else if (isa == ISA::RVV) { + using gemm_t = cpu_micro_gemm::MicroGemm; + if (has_zp) { + using dequantizer_t = Dequantizer4b; + cpu_gemm_wna16_impl( + input.data_ptr(), q_weight.data_ptr(), + output.data_ptr(), scales.data_ptr(), zeros_ptr, + g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_num, group_size, + pack_factor); + return; + } + if (use_desc_act) { + using dequantizer_t = Dequantizer4b; + cpu_gemm_wna16_impl( + input.data_ptr(), q_weight.data_ptr(), + output.data_ptr(), scales.data_ptr(), zeros_ptr, + g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_num, group_size, + pack_factor); + return; + } else { + using dequantizer_t = Dequantizer4b; + cpu_gemm_wna16_impl( + input.data_ptr(), q_weight.data_ptr(), + output.data_ptr(), scales.data_ptr(), zeros_ptr, + g_idx_ptr, bias.has_value() ? bias->data_ptr() : nullptr, + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_num, group_size, + pack_factor); + return; + } } }); } diff --git a/csrc/cpu/dnnl_kernels.cpp b/csrc/cpu/dnnl_kernels.cpp index 058fe25b0e2..6dda0929616 100644 --- a/csrc/cpu/dnnl_kernels.cpp +++ b/csrc/cpu/dnnl_kernels.cpp @@ -202,7 +202,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output, using cvt_vec_t = typename KernelVecType::cvt_vec_type; constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM; - const int64_t thread_num = omp_get_max_threads(); + const int64_t thread_num = cpu_utils::get_max_threads(); if (num_tokens > thread_num) { #pragma omp parallel for for (int64_t i = 0; i < num_tokens; ++i) { diff --git a/csrc/cpu/float_convert.hpp b/csrc/cpu/float_convert.hpp index c792bf131cc..0682ef40283 100644 --- a/csrc/cpu/float_convert.hpp +++ b/csrc/cpu/float_convert.hpp @@ -1,14 +1,15 @@ +#pragma once -static float bf16_to_float(uint16_t bf16) { +#include +#include + +inline float bf16_to_float(uint16_t bf16) { uint32_t bits = static_cast(bf16) << 16; - float fp32; - std::memcpy(&fp32, &bits, sizeof(fp32)); - return fp32; + return std::bit_cast(bits); } -static uint16_t float_to_bf16(float fp32) { - uint32_t bits; - std::memcpy(&bits, &fp32, sizeof(fp32)); +inline uint16_t float_to_bf16(float fp32) { + uint32_t bits = std::bit_cast(fp32); return static_cast(bits >> 16); } @@ -18,14 +19,13 @@ static uint16_t float_to_bf16(float fp32) { * Codes below copied from * https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab *************************************************/ -static uint16_t float_to_fp16(float fp32) { +inline uint16_t float_to_fp16(float fp32) { uint16_t fp16; - unsigned x; unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1; unsigned sign, exponent, mantissa; - std::memcpy(&x, &fp32, sizeof(fp32)); + uint32_t x = std::bit_cast(fp32); u = (x & 0x7fffffff); // Get rid of +NaN/-NaN case first. @@ -77,12 +77,11 @@ static uint16_t float_to_fp16(float fp32) { return fp16; } -static float fp16_to_float(uint16_t fp16) { +inline float fp16_to_float(uint16_t fp16) { unsigned sign = ((fp16 >> 15) & 1); unsigned exponent = ((fp16 >> 10) & 0x1f); unsigned mantissa = ((fp16 & 0x3ff) << 13); - int temp; - float fp32; + uint32_t temp; if (exponent == 0x1f) { /* NaN or Inf */ mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0); exponent = 0xff; @@ -101,6 +100,5 @@ static float fp16_to_float(uint16_t fp16) { exponent += 0x70; } temp = ((sign << 31) | (exponent << 23) | mantissa); - std::memcpy(&fp32, &temp, sizeof(temp)); - return fp32; + return std::bit_cast(temp); } diff --git a/csrc/cpu/generate_cpu_attn_dispatch.py b/csrc/cpu/generate_cpu_attn_dispatch.py index 7c7123a6def..95ce9e66927 100644 --- a/csrc/cpu/generate_cpu_attn_dispatch.py +++ b/csrc/cpu/generate_cpu_attn_dispatch.py @@ -11,7 +11,7 @@ import os HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512] # Head dimensions divisible by 16 but not 32 (VEC16 only) -HEAD_DIMS_16 = [80, 112] +HEAD_DIMS_16 = [48, 80, 112] # ISA types ISA_TYPES = { diff --git a/csrc/cpu/layernorm.cpp b/csrc/cpu/layernorm.cpp index a76ad08928a..704fb146338 100644 --- a/csrc/cpu/layernorm.cpp +++ b/csrc/cpu/layernorm.cpp @@ -4,8 +4,9 @@ namespace { template void rms_norm_impl(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, - const scalar_t* __restrict__ weight, const float epsilon, - const int num_tokens, const int hidden_size) { + const scalar_t* __restrict__ weight, const bool has_weight, + const float epsilon, const int num_tokens, + const int hidden_size) { using scalar_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0); @@ -27,12 +28,15 @@ void rms_norm_impl(scalar_t* __restrict__ out, for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) { scalar_vec_t x(input_p + j); - scalar_vec_t w(weight + j); - vec_op::FP32Vec8 fp32_x(x); - vec_op::FP32Vec8 fp32_w(w); - - vec_op::FP32Vec8 fp32_out = fp32_x * fp32_s_variance * fp32_w; + vec_op::FP32Vec8 fp32_out; + if (has_weight) { + scalar_vec_t w(weight + j); + vec_op::FP32Vec8 fp32_w(w); + fp32_out = fp32_x * fp32_s_variance * fp32_w; + } else { + fp32_out = fp32_x * fp32_s_variance; + } scalar_vec_t out(fp32_out); out.save(output_p + j); @@ -44,8 +48,8 @@ template void fused_add_rms_norm_impl(scalar_t* __restrict__ input, scalar_t* __restrict__ residual, const scalar_t* __restrict__ weight, - const float epsilon, const int num_tokens, - const int hidden_size) { + const bool has_weight, const float epsilon, + const int num_tokens, const int hidden_size) { using scalar_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0); @@ -72,13 +76,18 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input, vec_op::FP32Vec8 fp32_s_variance(s_variance); for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) { - scalar_vec_t w(weight + j); - scalar_vec_t res(residual_p + j); - - vec_op::FP32Vec8 fp32_w(w); - vec_op::FP32Vec8 fp32_res(res); - - vec_op::FP32Vec8 fp32_out = fp32_res * fp32_s_variance * fp32_w; + vec_op::FP32Vec8 fp32_out; + if (has_weight) { + scalar_vec_t w(weight + j); + scalar_vec_t res(residual_p + j); + vec_op::FP32Vec8 fp32_w(w); + vec_op::FP32Vec8 fp32_res(res); + fp32_out = fp32_res * fp32_s_variance * fp32_w; + } else { + scalar_vec_t res(residual_p + j); + vec_op::FP32Vec8 fp32_res(res); + fp32_out = fp32_res * fp32_s_variance; + } scalar_vec_t out(fp32_out); out.save(input_p + j); @@ -87,31 +96,41 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input, } } // namespace -void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, - double epsilon) { +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon) { int hidden_size = input.size(-1); int num_tokens = input.numel() / hidden_size; + const bool has_weight = weight.has_value(); + if (has_weight) { + TORCH_CHECK(weight->is_contiguous()); + } VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] { CPU_KERNEL_GUARD_IN(rms_norm_impl) rms_norm_impl(out.data_ptr(), input.data_ptr(), - weight.data_ptr(), epsilon, num_tokens, - hidden_size); + has_weight ? weight->data_ptr() : nullptr, + has_weight, epsilon, num_tokens, hidden_size); CPU_KERNEL_GUARD_OUT(rms_norm_impl) }); } void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - torch::Tensor& weight, double epsilon) { + std::optional weight, double epsilon) { int hidden_size = input.size(-1); int num_tokens = input.numel() / hidden_size; + const bool has_weight = weight.has_value(); + if (has_weight) { + TORCH_CHECK(weight->scalar_type() == input.scalar_type()); + TORCH_CHECK(weight->is_contiguous()); + } VLLM_DISPATCH_FLOATING_TYPES( input.scalar_type(), "fused_add_rms_norm_impl", [&] { CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl) fused_add_rms_norm_impl( input.data_ptr(), residual.data_ptr(), - weight.data_ptr(), epsilon, num_tokens, hidden_size); + has_weight ? weight->data_ptr() : nullptr, has_weight, + epsilon, num_tokens, hidden_size); CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl) }); } diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp index 357c7cf1d78..99e7c4a1d5c 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp @@ -213,6 +213,8 @@ class MicroGemm { public: static constexpr int32_t MaxMSize = 32; static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = 16; + static constexpr bool PackA = false; public: MicroGemm() : curr_m_(-1) { diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp index 23e78a681b5..f0471f71470 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp @@ -21,6 +21,9 @@ class MicroGemm { public: static constexpr int32_t MaxMSize = 16; static constexpr int32_t NSize = 16; + static constexpr int32_t WeightOCGroupSize = 16; + // callers must pack A matrix before GEMM + static constexpr bool PackA = false; public: void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp new file mode 100644 index 00000000000..7d4898852bb --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp @@ -0,0 +1,503 @@ +#ifndef CPU_MICRO_GEMM_NEON_HPP +#define CPU_MICRO_GEMM_NEON_HPP + +#include +#include + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#include +#include + +namespace cpu_micro_gemm { + +namespace { + +constexpr int32_t K = 4; +constexpr int32_t Cols = 2; +constexpr int32_t TileSize = K * Cols; +constexpr int32_t Mr = 8; +constexpr int32_t Nr = 8; +constexpr int32_t Nr_gemv = 16; + +// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1] +FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); +} + +// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a2, a3, b2, b3] +FORCE_INLINE float32x4_t zip2_f32x4(const float32x4_t a, const float32x4_t b) { + return vreinterpretq_f32_f64( + vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b))); +} + +FORCE_INLINE void init_acc_rowpair(float32x4_t& acc01, float32x4_t& acc23, + float32x4_t& acc45, float32x4_t& acc67, + const float* __restrict__ c_ptr, + const int64_t ldc, const int32_t m_rows, + const bool accum_c) { + if (!accum_c || m_rows == 0) { + acc01 = vdupq_n_f32(0.0f); + acc23 = vdupq_n_f32(0.0f); + acc45 = vdupq_n_f32(0.0f); + acc67 = vdupq_n_f32(0.0f); + return; + } + + const float32x4_t row0_0123 = vld1q_f32(c_ptr); + const float32x4_t row0_4567 = vld1q_f32(c_ptr + 4); + const float32x4_t row1_0123 = + (m_rows == 2) ? vld1q_f32(c_ptr + ldc) : vdupq_n_f32(0.0f); + const float32x4_t row1_4567 = + (m_rows == 2) ? vld1q_f32(c_ptr + ldc + 4) : vdupq_n_f32(0.0f); + + acc01 = zip1_f32x4(row0_0123, row1_0123); + acc23 = zip2_f32x4(row0_0123, row1_0123); + acc45 = zip1_f32x4(row0_4567, row1_4567); + acc67 = zip2_f32x4(row0_4567, row1_4567); +} + +FORCE_INLINE void store_acc_rowpair(const float32x4_t acc01, + const float32x4_t acc23, + const float32x4_t acc45, + const float32x4_t acc67, + float* __restrict__ c_ptr, + const int64_t ldc, const int32_t m_rows) { + if (m_rows == 0) { + return; + } + + vst1q_f32(c_ptr, zip1_f32x4(acc01, acc23)); + vst1q_f32(c_ptr + 4, zip1_f32x4(acc45, acc67)); + + if (m_rows == 2) { + vst1q_f32(c_ptr + ldc, zip2_f32x4(acc01, acc23)); + vst1q_f32(c_ptr + ldc + 4, zip2_f32x4(acc45, acc67)); + } +} + +FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a( + const bfloat16_t* __restrict__ a_packed, + const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr, + const int32_t m, const int32_t k_size, const int64_t ldc, + const bool accum_c) { + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc4501, acc4523, acc4545, acc4567; + float32x4_t acc6701, acc6723, acc6745, acc6767; + + init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, + std::min(2, m), accum_c); + init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + std::min(2, std::max(0, m - 2)), accum_c); + init_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc, + std::min(2, std::max(0, m - 4)), accum_c); + init_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc, + std::min(2, std::max(0, m - 6)), accum_c); + + const bfloat16_t* __restrict__ a_tile = a_packed; + const bfloat16_t* __restrict__ b_tile = b_packed; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile); + const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize); + const bfloat16x8_t a_tile45 = vld1q_bf16(a_tile + 2 * TileSize); + const bfloat16x8_t a_tile67 = vld1q_bf16(a_tile + 3 * TileSize); + + const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile); + const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile + TileSize); + const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile + 2 * TileSize); + const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile + 3 * TileSize); + + acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01); + acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01); + acc4501 = vbfmmlaq_f32(acc4501, a_tile45, b_tile01); + acc6701 = vbfmmlaq_f32(acc6701, a_tile67, b_tile01); + + acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23); + acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23); + acc4523 = vbfmmlaq_f32(acc4523, a_tile45, b_tile23); + acc6723 = vbfmmlaq_f32(acc6723, a_tile67, b_tile23); + + acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45); + acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45); + acc4545 = vbfmmlaq_f32(acc4545, a_tile45, b_tile45); + acc6745 = vbfmmlaq_f32(acc6745, a_tile67, b_tile45); + + acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67); + acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67); + acc4567 = vbfmmlaq_f32(acc4567, a_tile45, b_tile67); + acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67); + + a_tile += 4 * TileSize; + b_tile += Nr * K; + } + + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, + std::min(2, m)); + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + std::min(2, std::max(0, m - 2))); + store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc, + std::min(2, std::max(0, m - 4))); + store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc, + std::min(2, std::max(0, m - 6))); +} + +FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a( + const bfloat16_t* __restrict__ a_packed, + const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr, + const int32_t m, const int32_t k_size, const int64_t b_n_group_stride, + const int64_t ldc, const bool accum_c) { + const int32_t m_rows_01 = std::min(2, m); + const int32_t m_rows_23 = std::min(2, std::max(0, m - 2)); + + float32x4_t acc0101, acc0123, acc0145, acc0167; + float32x4_t acc2301, acc2323, acc2345, acc2367; + float32x4_t acc0189, acc011011, acc011213, acc011415; + float32x4_t acc2389, acc231011, acc231213, acc231415; + + init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01, + accum_c); + init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + m_rows_23, accum_c); + init_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc, + m_rows_01, accum_c); + init_acc_rowpair(acc2389, acc231011, acc231213, acc231415, + c_ptr + 2 * ldc + 8, ldc, m_rows_23, accum_c); + + const bfloat16_t* __restrict__ a_tile = a_packed; + const bfloat16_t* __restrict__ b_tile0 = b_packed; + const bfloat16_t* __restrict__ b_tile1 = b_packed + b_n_group_stride; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile); + const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize); + const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile0); + const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile0 + TileSize); + const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile0 + 2 * TileSize); + const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile0 + 3 * TileSize); + const bfloat16x8_t b_tile89 = vld1q_bf16(b_tile1); + const bfloat16x8_t b_tile1011 = vld1q_bf16(b_tile1 + TileSize); + const bfloat16x8_t b_tile1213 = vld1q_bf16(b_tile1 + 2 * TileSize); + const bfloat16x8_t b_tile1415 = vld1q_bf16(b_tile1 + 3 * TileSize); + + acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01); + acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01); + acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23); + acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23); + + acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45); + acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45); + acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67); + acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67); + + acc0189 = vbfmmlaq_f32(acc0189, a_tile01, b_tile89); + acc2389 = vbfmmlaq_f32(acc2389, a_tile23, b_tile89); + acc011011 = vbfmmlaq_f32(acc011011, a_tile01, b_tile1011); + acc231011 = vbfmmlaq_f32(acc231011, a_tile23, b_tile1011); + + acc011213 = vbfmmlaq_f32(acc011213, a_tile01, b_tile1213); + acc231213 = vbfmmlaq_f32(acc231213, a_tile23, b_tile1213); + acc011415 = vbfmmlaq_f32(acc011415, a_tile01, b_tile1415); + acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415); + + a_tile += 2 * TileSize; + b_tile0 += Nr * K; + b_tile1 += Nr * K; + } + + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01); + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + m_rows_23); + store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc, + m_rows_01); + store_acc_rowpair(acc2389, acc231011, acc231213, acc231415, + c_ptr + 2 * ldc + 8, ldc, m_rows_23); +} + +} // namespace + +template +class MicroGemm { + public: + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = Nr; + static constexpr bool PackA = false; + + public: + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { + TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16."); + } + + static void pack_weight(const scalar_t* __restrict__ /*weight*/, + scalar_t* __restrict__ /*packed_weight*/, + const int32_t /*output_size*/, + const int32_t /*input_size*/) { + TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16."); + } +}; + +template <> +class MicroGemm { + public: + using scalar_t = c10::BFloat16; + + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = Nr; + static constexpr bool PackA = true; + + public: + // physical layout [ + // M / 8; Mr is 8 + // K / 4; K for bfmmla is 4 + // 4, ; 4 row-pairs for each 8 rows + // 2, ; row-pair is 2 rows + // 4 ; 4 elements per row + // ] + + static void pack_input_from_rows(const scalar_t* const* __restrict__ rows, + scalar_t* __restrict__ a_packed, + const int32_t m, const int32_t k) { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK_EQ(k % K, 0); + + auto* __restrict__ out = reinterpret_cast(a_packed); + const bfloat16x8_t zero_q = vdupq_n_bf16(bfloat16_t{}); + const bfloat16x4_t zero = vget_low_bf16(zero_q); + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t actual_m = std::min(Mr, m - row_base); + const bfloat16_t* __restrict__ row[Mr]; + for (int32_t i = 0; i < actual_m; ++i) { + row[i] = reinterpret_cast(rows[row_base + i]); + } + + if (actual_m == 8) { + int32_t k_idx = 0; + for (; k_idx + 8 <= k; k_idx += 8) { + bfloat16_t* __restrict__ block0 = out; + bfloat16_t* __restrict__ block1 = out + 4 * TileSize; + + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx); + bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx); + vst1q_bf16(block0, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = vld1q_bf16(row[2] + k_idx); + a1 = vld1q_bf16(row[3] + k_idx); + vst1q_bf16(block0 + TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = vld1q_bf16(row[4] + k_idx); + a1 = vld1q_bf16(row[5] + k_idx); + vst1q_bf16(block0 + 2 * TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + 2 * TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = vld1q_bf16(row[6] + k_idx); + a1 = vld1q_bf16(row[7] + k_idx); + vst1q_bf16(block0 + 3 * TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + 3 * TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + out += 8 * TileSize; + } + + for (; k_idx < k; k_idx += K) { + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx); + bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx); + vst1q_bf16(out, vcombine_bf16(a0, a1)); + + a0 = vld1_bf16(row[2] + k_idx); + a1 = vld1_bf16(row[3] + k_idx); + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1)); + + a0 = vld1_bf16(row[4] + k_idx); + a1 = vld1_bf16(row[5] + k_idx); + vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1)); + + a0 = vld1_bf16(row[6] + k_idx); + a1 = vld1_bf16(row[7] + k_idx); + vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1)); + + out += 4 * TileSize; + } + continue; + } + + if (actual_m == 4) { + int32_t k_idx = 0; + for (; k_idx + 8 <= k; k_idx += 8) { + bfloat16_t* __restrict__ block0 = out; + bfloat16_t* __restrict__ block1 = out + 2 * TileSize; + + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx); + bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx); + vst1q_bf16(block0, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = vld1q_bf16(row[2] + k_idx); + a1 = vld1q_bf16(row[3] + k_idx); + vst1q_bf16(block0 + TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + out += 4 * TileSize; + } + + for (; k_idx < k; k_idx += K) { + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx); + bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx); + vst1q_bf16(out, vcombine_bf16(a0, a1)); + + a0 = vld1_bf16(row[2] + k_idx); + a1 = vld1_bf16(row[3] + k_idx); + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1)); + + out += 2 * TileSize; + } + continue; + } + + const int32_t row_pair_count = (actual_m <= 4) ? 2 : Mr / 2; + + int32_t k_idx = 0; + for (; k_idx + 8 <= k; k_idx += 8) { + bfloat16_t* __restrict__ block0 = out; + bfloat16_t* __restrict__ block1 = out + row_pair_count * TileSize; + + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx); + bfloat16x8_t a1 = (actual_m > 1) ? vld1q_bf16(row[1] + k_idx) : zero_q; + vst1q_bf16(block0, vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = (actual_m > 2) ? vld1q_bf16(row[2] + k_idx) : zero_q; + a1 = (actual_m > 3) ? vld1q_bf16(row[3] + k_idx) : zero_q; + vst1q_bf16(block0 + TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + if (actual_m > 4) { + a0 = vld1q_bf16(row[4] + k_idx); + a1 = (actual_m > 5) ? vld1q_bf16(row[5] + k_idx) : zero_q; + vst1q_bf16(block0 + 2 * TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + 2 * TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + + a0 = (actual_m > 6) ? vld1q_bf16(row[6] + k_idx) : zero_q; + a1 = (actual_m > 7) ? vld1q_bf16(row[7] + k_idx) : zero_q; + vst1q_bf16(block0 + 3 * TileSize, + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1))); + vst1q_bf16(block1 + 3 * TileSize, + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1))); + } + + out += 2 * row_pair_count * TileSize; + } + + for (; k_idx < k; k_idx += K) { + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx); + bfloat16x4_t a1 = (actual_m > 1) ? vld1_bf16(row[1] + k_idx) : zero; + vst1q_bf16(out, vcombine_bf16(a0, a1)); + + a0 = (actual_m > 2) ? vld1_bf16(row[2] + k_idx) : zero; + a1 = (actual_m > 3) ? vld1_bf16(row[3] + k_idx) : zero; + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1)); + + if (actual_m > 4) { + a0 = vld1_bf16(row[4] + k_idx); + a1 = (actual_m > 5) ? vld1_bf16(row[5] + k_idx) : zero; + vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1)); + + a0 = (actual_m > 6) ? vld1_bf16(row[6] + k_idx) : zero; + a1 = (actual_m > 7) ? vld1_bf16(row[7] + k_idx) : zero; + vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1)); + } + out += row_pair_count * TileSize; + } + } + } + + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { + (void)lda; // A is packed, so lda is not needed + TORCH_CHECK_EQ(k % K, 0); + + for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) { + const bfloat16_t* __restrict__ b_panel = + reinterpret_cast(b_ptr) + n_idx * k; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const bfloat16_t* __restrict__ a_panel = + reinterpret_cast(a_ptr) + row_base * k; + float* __restrict__ c_panel = c_ptr + row_base * ldc + n_idx; + + if (panel_m <= 4) { + gemm_micro_bfmmla_4x16_packed_a(a_panel, b_panel, c_panel, panel_m, k, + b_n_group_stride, ldc, accum_c); + } else { + gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel, c_panel, panel_m, k, + ldc, accum_c); + gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel + b_n_group_stride, + c_panel + Nr, panel_m, k, ldc, + accum_c); + } + } + } + } + + // physical layout [ + // N / 8; Nr is 8 + // K / 4; K for bfmmla is 4 + // 4, ; 4 col-pairs for each 8 cols + // 2, ; col-pair is 2 cols + // 4 ; 4 elements per col + // ] + static void pack_weight(const c10::BFloat16* __restrict__ weight, + c10::BFloat16* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK_EQ(output_size % NSize, 0); + TORCH_CHECK_EQ(input_size % K, 0); + + for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) { + c10::BFloat16* __restrict__ dst = packed_weight + o_idx * input_size; + for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) { + for (int32_t pair_idx = 0; pair_idx < Nr; pair_idx += Cols) { + const c10::BFloat16* __restrict__ row0 = + weight + (o_idx + pair_idx) * input_size; + const c10::BFloat16* __restrict__ row1 = row0 + input_size; + dst[0] = row0[k_idx + 0]; + dst[1] = row0[k_idx + 1]; + dst[2] = row0[k_idx + 2]; + dst[3] = row0[k_idx + 3]; + dst[4] = row1[k_idx + 0]; + dst[5] = row1[k_idx + 1]; + dst[6] = row1[k_idx + 2]; + dst[7] = row1[k_idx + 3]; + dst += TileSize; + } + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_rvv.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_rvv.hpp new file mode 100644 index 00000000000..3e3c056f649 --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_rvv.hpp @@ -0,0 +1,228 @@ +#ifndef CPU_MICRO_GEMM_RVV_HPP +#define CPU_MICRO_GEMM_RVV_HPP + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#if defined(__riscv_v) + +namespace cpu_micro_gemm { +namespace { + +constexpr int32_t RVV_MGEMM_N8 = 8; +constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16; + +template +FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr); + +template <> +FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const float* ptr) { + return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8); +} + +template <> +FORCE_INLINE fixed_fp32x8_t +load_row8_b_as_f32(const c10::Half* ptr) { + #if defined(__riscv_zvfh) + fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)( + reinterpret_cast(ptr), RVV_MGEMM_N8); + return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8); + #else + alignas(32) float values[RVV_MGEMM_N8]; + for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) { + values[i] = static_cast(ptr[i]); + } + return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8); + #endif +} + +template <> +FORCE_INLINE fixed_fp32x8_t +load_row8_b_as_f32(const c10::BFloat16* ptr) { + #if defined(__riscv_zvfbfmin) + fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)( + reinterpret_cast(ptr), RVV_MGEMM_N8); + fixed_bf16x8_t vec = + RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw); + return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8); + #else + fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)( + reinterpret_cast(ptr), RVV_MGEMM_N8); + auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8); + auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8); + return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted); + #endif +} + +// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group, +// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements. +template +FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr, + const scalar_t* __restrict__ b_ptr, + float* __restrict__ c_ptr, + const int64_t lda, + const int64_t ldc, const int32_t k, + const bool accum_c) { + static_assert(0 < M && M <= 8); + + #define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7) + #define RVV_IF_M(i) if constexpr (M > (i)) + + #define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda; + RVV_ROWS_APPLY(RVV_DECL_A) + #undef RVV_DECL_A + + #define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i; + RVV_ROWS_APPLY(RVV_DECL_ACC) + #undef RVV_DECL_ACC + + #define RVV_INIT_ACC(i) \ + RVV_IF_M(i) { \ + if (accum_c) { \ + acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \ + RVV_MGEMM_N8); \ + } else { \ + acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \ + } \ + } + RVV_ROWS_APPLY(RVV_INIT_ACC) + #undef RVV_INIT_ACC + + int32_t k_idx = 0; + for (; k_idx + 3 < k; k_idx += 4) { + #define RVV_FMA_ROW(i, K_OFFSET) \ + RVV_IF_M(i) { \ + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \ + acc##i, static_cast(*(a##i + k_idx + (K_OFFSET))), b, \ + RVV_MGEMM_N8); \ + } + + #define RVV_STEP_K(K_OFFSET) \ + { \ + fixed_fp32x8_t b = load_row8_b_as_f32( \ + b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \ + RVV_FMA_ROW(0, K_OFFSET) \ + RVV_FMA_ROW(1, K_OFFSET) \ + RVV_FMA_ROW(2, K_OFFSET) \ + RVV_FMA_ROW(3, K_OFFSET) \ + RVV_FMA_ROW(4, K_OFFSET) \ + RVV_FMA_ROW(5, K_OFFSET) \ + RVV_FMA_ROW(6, K_OFFSET) \ + RVV_FMA_ROW(7, K_OFFSET) \ + } + + RVV_STEP_K(0) + RVV_STEP_K(1) + RVV_STEP_K(2) + RVV_STEP_K(3) + #undef RVV_STEP_K + #undef RVV_FMA_ROW + } + + for (; k_idx < k; ++k_idx) { + fixed_fp32x8_t b = + load_row8_b_as_f32(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE); + #define RVV_TAIL_ROW(i) \ + RVV_IF_M(i) { \ + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \ + acc##i, static_cast(*(a##i + k_idx)), b, RVV_MGEMM_N8); \ + } + RVV_ROWS_APPLY(RVV_TAIL_ROW) + #undef RVV_TAIL_ROW + } + + #define RVV_STORE_ROW(i) \ + RVV_IF_M(i) { \ + RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \ + RVV_MGEMM_N8); \ + } + RVV_ROWS_APPLY(RVV_STORE_ROW) + #undef RVV_STORE_ROW + + #undef RVV_ROWS_APPLY + #undef RVV_IF_M +} + +template +FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) { + static_assert(0 < M && M <= 8); + scalar_t* __restrict__ curr_b_0 = b_ptr; + scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride; + + gemm_micro_rvv_fma_mx8_ku4(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c); + gemm_micro_rvv_fma_mx8_ku4(a_ptr, curr_b_0 + RVV_MGEMM_N8, + c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c); + gemm_micro_rvv_fma_mx8_ku4(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k, + accum_c); + gemm_micro_rvv_fma_mx8_ku4(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda, + ldc, k, accum_c); +} + +class TileGemmRVV { + public: + template + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { + switch (m) { + case 1: + gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS); + break; + case 2: + gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS); + break; + case 3: + gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS); + break; + case 4: + gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS); + break; + case 5: + gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS); + break; + case 6: + gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS); + break; + case 7: + gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS); + break; + case 8: + gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS); + break; + } + } +}; + +} // namespace + +template +class MicroGemm { + public: + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = 32; + + public: + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { + TileGemmRVV::gemm(CPU_MICRO_GEMM_PARAMS); + } + + static void pack_weight(const scalar_t* __restrict__ weight, + scalar_t* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK_EQ(output_size % 16, 0); + for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) { + const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size; + scalar_t* __restrict__ curr_packed_weight = + packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16; + for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) { + *curr_packed_weight = *curr_weight; + + curr_packed_weight += 16; + ++curr_weight; + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif // defined(__riscv_v) + +#endif // CPU_MICRO_GEMM_RVV_HPP diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp index 1c605a2851d..ad7d4be113e 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp @@ -104,6 +104,8 @@ class MicroGemm { public: static constexpr int32_t MaxMSize = 8; static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = 16; + static constexpr bool PackA = false; public: void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) { diff --git a/csrc/cpu/mla_decode.cpp b/csrc/cpu/mla_decode.cpp index 582c480c3be..702912a5bcc 100644 --- a/csrc/cpu/mla_decode.cpp +++ b/csrc/cpu/mla_decode.cpp @@ -18,17 +18,9 @@ struct KernelVecType { template <> struct KernelVecType { -#if defined(__powerpc64__) - // Power specific vector types - using qk_load_vec_type = vec_op::FP32Vec16; - using qk_vec_type = vec_op::FP32Vec16; - using v_load_vec_type = vec_op::FP32Vec16; -#else - // Fallback for other architectures, including x86 using qk_load_vec_type = vec_op::FP16Vec16; using qk_vec_type = vec_op::FP32Vec16; using v_load_vec_type = vec_op::FP16Vec16; -#endif }; #ifdef __AVX512BF16__ @@ -259,7 +251,7 @@ void mla_decode_kvcache_cpu_impl( constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM; // shared across threads - const int max_threads = omp_get_max_threads(); + const int max_threads = cpu_utils::get_max_threads(); const int acc_out_nbytes = max_threads * num_heads * V_HEAD_DIM * sizeof(float); float* acc_out = static_cast(std::aligned_alloc(64, acc_out_nbytes)); diff --git a/csrc/cpu/pos_encoding.cpp b/csrc/cpu/pos_encoding.cpp index 9f41e4e222b..b241918902e 100644 --- a/csrc/cpu/pos_encoding.cpp +++ b/csrc/cpu/pos_encoding.cpp @@ -1,4 +1,3 @@ - #include "cpu_types.hpp" namespace { @@ -97,6 +96,91 @@ void rotary_embedding_impl( } } +template <> +void rotary_embedding_impl( + const int64_t* __restrict__ positions, c10::Half* __restrict__ query, + c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache, + const int rot_dim, const int64_t query_stride, const int64_t key_stride, + const int num_heads, const int num_kv_heads, const int head_size, + const int num_tokens) { + using scalar_vec_t = vec_op::FP16Vec8; + constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); + + const int embed_dim = rot_dim / 2; + bool flag = (embed_dim % VEC_ELEM_NUM == 0); + const int loop_upper = flag ? embed_dim : embed_dim - VEC_ELEM_NUM; + + auto compute_loop = [&](const int64_t token_head, const c10::Half* cache_ptr, + c10::Half* qk) { + int j = 0; + for (; j < loop_upper; j += VEC_ELEM_NUM) { + const int rot_offset = j; + const int x_index = rot_offset; + const int y_index = embed_dim + rot_offset; + + const int64_t out_x = token_head + x_index; + const int64_t out_y = token_head + y_index; + + const vec_op::FP16Vec8 cos_fp16(cache_ptr + x_index); + const vec_op::FP16Vec8 sin_fp16(cache_ptr + y_index); + const vec_op::FP16Vec8 q_x_fp16(qk + out_x); + const vec_op::FP16Vec8 q_y_fp16(qk + out_y); + + const vec_op::FP32Vec8 fp32_cos(cos_fp16); + const vec_op::FP32Vec8 fp32_sin(sin_fp16); + const vec_op::FP32Vec8 fp32_q_x(q_x_fp16); + const vec_op::FP32Vec8 fp32_q_y(q_y_fp16); + + auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin; + auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin; + + vec_op::FP16Vec8(out1).save(qk + out_x); + vec_op::FP16Vec8(out2).save(qk + out_y); + } + if (!flag) { + for (; j < embed_dim; ++j) { + const int x_index = j; + const int y_index = embed_dim + j; + + const int64_t out_x = token_head + x_index; + const int64_t out_y = token_head + y_index; + + const float fp32_cos = static_cast(cache_ptr[x_index]); + const float fp32_sin = static_cast(cache_ptr[y_index]); + const float fp32_q_x = static_cast(qk[out_x]); + const float fp32_q_y = static_cast(qk[out_y]); + + qk[out_x] = + static_cast(fp32_q_x * fp32_cos - fp32_q_y * fp32_sin); + qk[out_y] = + static_cast(fp32_q_y * fp32_cos + fp32_q_x * fp32_sin); + } + } + }; + +#pragma omp parallel for + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + + for (int i = 0; i < num_heads; ++i) { + const int head_idx = i; + const int64_t token_head = + token_idx * query_stride + head_idx * head_size; + compute_loop(token_head, cache_ptr, query); + } + + if (key != nullptr) { + for (int i = 0; i < num_kv_heads; ++i) { + const int head_idx = i; + const int64_t token_head = + token_idx * key_stride + head_idx * head_size; + compute_loop(token_head, cache_ptr, key); + } + } + } +} + template void rotary_embedding_gptj_impl( const int64_t* __restrict__ positions, // [batch_size, seq_len] or @@ -174,6 +258,75 @@ void rotary_embedding_gptj_impl( } } } + +template <> +void rotary_embedding_gptj_impl( + const int64_t* __restrict__ positions, c10::Half* __restrict__ query, + c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache, + const int rot_dim, const int64_t query_stride, const int64_t key_stride, + const int num_heads, const int num_kv_heads, const int head_size, + const int num_tokens) { + const int embed_dim = rot_dim / 2; + +#pragma omp parallel for collapse(2) + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + for (int i = 0; i < num_heads; ++i) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + const c10::Half* cos_cache_ptr = cache_ptr; + const c10::Half* sin_cache_ptr = cache_ptr + embed_dim; + const int head_idx = i; + const int64_t token_head = + token_idx * query_stride + head_idx * head_size; + c10::Half* head_query = token_head + query; + for (int j = 0; j < embed_dim; j += 1) { + const int rot_offset = j; + const int x_index = 2 * rot_offset; + const int y_index = 2 * rot_offset + 1; + + const float cos = static_cast(cos_cache_ptr[rot_offset]); + const float sin = static_cast(sin_cache_ptr[rot_offset]); + + const float x = static_cast(head_query[x_index]); + const float y = static_cast(head_query[y_index]); + + head_query[x_index] = static_cast(x * cos - y * sin); + head_query[y_index] = static_cast(y * cos + x * sin); + } + } + } + + if (key == nullptr) { + return; + } + +#pragma omp parallel for collapse(2) + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + for (int i = 0; i < num_kv_heads; ++i) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + const c10::Half* cos_cache_ptr = cache_ptr; + const c10::Half* sin_cache_ptr = cache_ptr + embed_dim; + const int head_idx = i; + const int64_t token_head = token_idx * key_stride + head_idx * head_size; + c10::Half* head_key = key + token_head; + for (int j = 0; j < embed_dim; j += 1) { + const int rot_offset = j; + const int x_index = 2 * rot_offset; + const int y_index = 2 * rot_offset + 1; + + const float cos = static_cast(cos_cache_ptr[rot_offset]); + const float sin = static_cast(sin_cache_ptr[rot_offset]); + + const float x = static_cast(head_key[x_index]); + const float y = static_cast(head_key[y_index]); + + head_key[x_index] = static_cast(x * cos - y * sin); + head_key[y_index] = static_cast(y * cos + x * sin); + } + } + } +} }; // namespace void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, diff --git a/csrc/cpu/sgl-kernels/conv.cpp b/csrc/cpu/sgl-kernels/conv.cpp index 15114732aac..b918aed8bff 100644 --- a/csrc/cpu/sgl-kernels/conv.cpp +++ b/csrc/cpu/sgl-kernels/conv.cpp @@ -289,19 +289,18 @@ void causal_conv1d_fwd_kernel_impl( } } -#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \ - tinygemm_kernel::apply( \ - input + batch_offset * dim + mb_start * dim + nb_start, \ - weight + nb_start * width, \ - out + batch_offset * dim + mb_start * dim + nb_start, \ - has_bias ? bias + nb_start : nullptr, \ - nullptr, \ - false, \ - mb_size, \ - dim, \ +#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \ + tinygemm_kernel::apply( \ + input + batch_offset * dim + mb_start * dim + nb_start, \ + weight + nb_start * width, \ + out + batch_offset * dim + mb_start * dim + nb_start, \ + has_bias ? bias + nb_start : nullptr, \ + has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \ + has_initial_states_value, \ + mb_size, \ + dim, \ mb_start == 0); -// TODO: add `has_initial_state` support for varlen kernel template void causal_conv1d_fwd_varlen_kernel_impl( scalar_t* __restrict__ out, @@ -343,6 +342,9 @@ void causal_conv1d_fwd_varlen_kernel_impl( int64_t nb_start = nb * BLOCK_N; int64_t nb_size = std::min(dim - nb_start, BLOCK_N); + const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false; + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs; + switch (width << 4 | nb_size >> 4) { case 0x42: LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32); @@ -373,7 +375,7 @@ void causal_conv1d_fwd_varlen_kernel_impl( width, dim, seqlen, - /* has_initial_state */ false); + has_initial_state[bs]); } }); } diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index bf1b6444bdd..b9a793c599f 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -1116,6 +1116,164 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( }); } +// Speculative-decode variant: processes a varlen batch where each sequence has +// ``q_len`` draft tokens, runs the recurrence sequentially over those tokens +// (inside the kernel, so one dispatch handles the whole draft block), reads the +// initial state from cache slot ``num_accepted-1`` and stores the state *after* +// token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU +// kernel). Parallelized over (sequence, v_head); the per-sequence token loop is +// sequential as required by the recurrence. +template +void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl( + const scalar_t* __restrict__ q_ptr, // [T, HK, EK] + const scalar_t* __restrict__ k_ptr, // [T, HK, EK] + const scalar_t* __restrict__ v_ptr, // [T, HV, EV] + const param_t* __restrict__ A_log_ptr, + const scalar_t* __restrict__ a_ptr, // [T, HV] + const scalar_t* __restrict__ dt_bias_ptr, + const scalar_t* __restrict__ b_ptr, // [T, HV] + const int32_t* __restrict__ spec_indices_ptr, // [N, S] + const int32_t* __restrict__ num_accepted_ptr, // [N] + const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] + float* __restrict__ state_ptr, + scalar_t* __restrict__ o_ptr, // [T, HV, EV] + float* __restrict__ qk_scale_buf, // [2, T, HK] + int64_t total_tokens, + int64_t batch_size, + int64_t spec_stride, + int64_t num_heads, + int64_t head_dim, + int64_t v_num_heads, + int64_t v_head_dim, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH, + int64_t v_strideT, + int64_t v_strideH, + int64_t state_slot_stride, + bool use_qk_l2norm_in_kernel, + double softplus_threshold) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int64_t VecSize = bVec::size(); + constexpr int64_t fVecSize = fVec::size(); + int64_t group_size = v_num_heads / num_heads; + double scale = 1 / std::sqrt((double)head_dim); + fVec scale_vec = fVec((float)scale); + + if (use_qk_l2norm_in_kernel) { + float eps = 1e-5f; + at::parallel_for(0, total_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; ++i) { + int64_t ti = i / num_heads; + int64_t ni = i % num_heads; + const scalar_t* qp = q_ptr + ti * q_strideT + ni * q_strideH; + const scalar_t* kp = k_ptr + ti * k_strideT + ni * k_strideH; + float sq = 0.f, sk = 0.f; + for (int64_t d = 0; d < head_dim; ++d) { + float qv = (float)qp[d]; + sq += qv * qv; + float kv = (float)kp[d]; + sk += kv * kv; + } + qk_scale_buf[ti * num_heads + ni] = 1.f / std::sqrt(sq + eps); + qk_scale_buf[total_tokens * num_heads + ti * num_heads + ni] = 1.f / std::sqrt(sk + eps); + } + }); + } + + at::parallel_for(0, batch_size * v_num_heads, 0, [&](int64_t begin, int64_t end) { + for (int64_t idx = begin; idx < end; ++idx) { + int64_t bi = idx / v_num_heads; + int64_t ni = idx % v_num_heads; + int64_t kh = ni / group_size; + int64_t q_start = cu_seqlens_ptr[bi]; + int64_t q_len = cu_seqlens_ptr[bi + 1] - q_start; + if (q_len <= 0) { + continue; + } + int64_t acc = (int64_t)num_accepted_ptr[bi]; + // Clamp acc-1 to >=0: when num_accepted is 0 the unclamped index reads + // out of bounds and yields an arbitrary prev_slot used to index the SSM + // state. Mirrors the GPU guard tl.maximum(num_accepted - 1, 0). + int64_t prev_slot = + (int64_t)spec_indices_ptr[bi * spec_stride + (acc > 0 ? acc - 1 : 0)]; + for (int64_t t = 0; t < q_len; ++t) { + int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t]; + int64_t token = q_start + t; + const float* src = state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim; + float* dst = state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim; + float g_val = -std::exp((float)A_log_ptr[ni]) * + softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold); + float g_val_exp = std::exp(g_val); + fVec g_val_exp_vec = fVec(g_val_exp); + float beta_val = 1.f / (1.f + std::exp(-(float)b_ptr[token * v_num_heads + ni])); + fVec beta_vec = fVec(beta_val); + int64_t q_offset = token * q_strideT + kh * q_strideH; + int64_t k_offset = token * k_strideT + kh * k_strideH; + float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[token * num_heads + kh] : 1.f; + float k_scale = + use_qk_l2norm_in_kernel ? qk_scale_buf[total_tokens * num_heads + token * num_heads + kh] : 1.f; + int64_t v_offset = token * v_strideT + ni * v_strideH; + int64_t o_offset = (token * v_num_heads + ni) * v_head_dim; + int64_t dvi = 0; + for (; dvi <= v_head_dim - VecSize; dvi += VecSize) { + fVec kv_mem_vec0 = fVec(0.f); + fVec kv_mem_vec1 = fVec(0.f); + for (int di = 0; di < head_dim; ++di) { + fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale); + fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi); + fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize); + kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec; + kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec; + } + bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi); + fVec v_vec0, v_vec1; + std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec); + fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec; + fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec; + fVec o_vec0 = fVec(0.f); + fVec o_vec1 = fVec(0.f); + for (int di = 0; di < head_dim; ++di) { + fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale); + fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale); + fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi); + fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize); + sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0; + sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1; + o_vec0 = o_vec0 + sv0 * q_vec * scale_vec; + o_vec1 = o_vec1 + sv1 * q_vec * scale_vec; + sv0.store(dst + di * v_head_dim + dvi); + sv1.store(dst + di * v_head_dim + dvi + fVecSize); + } + bVec o_vec = at::vec::convert_from_float(o_vec0, o_vec1); + o_vec.store(o_ptr + o_offset + dvi); + } + for (; dvi < v_head_dim; ++dvi) { + float kv_mem_val = 0.f; + for (int di = 0; di < head_dim; ++di) { + float k_val = (float)k_ptr[k_offset + di] * k_scale; + kv_mem_val += src[di * v_head_dim + dvi] * g_val_exp * k_val; + } + float v_val = (float)v_ptr[v_offset + dvi]; + float dt_val = (v_val - kv_mem_val) * beta_val; + float o_val = 0.f; + for (int di = 0; di < head_dim; ++di) { + float q_val = (float)q_ptr[q_offset + di] * q_scale; + float k_val = (float)k_ptr[k_offset + di] * k_scale; + float ns = src[di * v_head_dim + dvi] * g_val_exp + k_val * dt_val; + dst[di * v_head_dim + dvi] = ns; + o_val += ns * q_val * scale; + } + o_ptr[o_offset + dvi] = (scalar_t)o_val; + } + prev_slot = cur_slot; + } + } + }); +} + template void fused_gdn_gating_kernel_impl( float* __restrict__ A_log, @@ -1500,6 +1658,103 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( return core_attn_out; } +// Speculative-decode update (multi-token, multi-slot rollback). +// q: [T, HK, EK] k: [T, HK, EK] v: [T, HV, EV] +// a: [T, HV] b: [T, HV] +// initial_state_source: [N_slots, HV, EK, EV] FP32 (updated in place) +// spec_state_indices: [batch, S] INT32 (S = num_spec + 1) +// num_accepted_tokens: [batch] INT32 +// cu_seqlens: [batch + 1] INT32 +// Returns output: [T, HV, EV] +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + at::Tensor& initial_state_source, + const at::Tensor& spec_state_indices, + const at::Tensor& num_accepted_tokens, + const at::Tensor& cu_seqlens, + bool use_qk_l2norm_in_kernel, + double softplus_beta = 1.0, + double softplus_threshold = 20.0) { + CHECK_DIM(3, q); + CHECK_DIM(3, v); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); + int64_t total_tokens = q.size(0); + int64_t num_heads = q.size(1); + int64_t head_dim = q.size(2); + int64_t v_num_heads = v.size(1); + int64_t v_head_dim = v.size(2); + int64_t batch_size = cu_seqlens.size(0) - 1; + int64_t spec_stride = spec_state_indices.stride(0); + CHECK_INPUT_SHAPE_DTYPE(k, {total_tokens, num_heads, head_dim}, q.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(a, {total_tokens, v_num_heads}, q.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(b, {total_tokens, v_num_heads}, q.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(dt_bias, {v_num_heads}, q.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(num_accepted_tokens, {batch_size}, at::kInt); + CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {batch_size + 1}, at::kInt); + CHECK_EQ(v_num_heads % num_heads, 0); + TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads})); + CHECK_INPUT_SHAPE_DTYPE( + initial_state_source, + {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}, + at::kFloat); + TORCH_CHECK(initial_state_source.size(0) >= batch_size, + "initial_state_source capacity too small: size(0)=", + initial_state_source.size(0), ", batch_size=", batch_size); + + int64_t q_strideT = q.stride(0); + int64_t q_strideH = q.stride(1); + int64_t k_strideT = k.stride(0); + int64_t k_strideH = k.stride(1); + int64_t v_strideT = v.stride(0); + int64_t v_strideH = v.stride(1); + int64_t state_slot_stride = initial_state_source.stride(0); + + at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options()); + at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat); + + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT( + q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] { + fused_sigmoid_gating_delta_rule_update_spec_kernel_impl( + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + A_log.data_ptr(), + a.data_ptr(), + dt_bias.data_ptr(), + b.data_ptr(), + spec_state_indices.data_ptr(), + num_accepted_tokens.data_ptr(), + cu_seqlens.data_ptr(), + initial_state_source.data_ptr(), + o.data_ptr(), + qk_scale_buf.data_ptr(), + total_tokens, + batch_size, + spec_stride, + num_heads, + head_dim, + v_num_heads, + v_head_dim, + q_strideT, + q_strideH, + k_strideT, + k_strideH, + v_strideT, + v_strideH, + state_slot_stride, + use_qk_l2norm_in_kernel, + softplus_threshold); + }); + return o; +} + // A_log: [num_v_heads] // a: [batch, num_v_heads] // b: [batch, num_v_heads] diff --git a/csrc/cpu/sgl-kernels/gemm_int4.cpp b/csrc/cpu/sgl-kernels/gemm_int4.cpp index 5b66b2a5aee..6dbd09080d0 100644 --- a/csrc/cpu/sgl-kernels/gemm_int4.cpp +++ b/csrc/cpu/sgl-kernels/gemm_int4.cpp @@ -268,6 +268,142 @@ void _dequant_gemm_accum_small_M( _dequant_gemm_accum_small_M(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc); #endif +template +inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) { + // B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two + // columns from adjacent 8-column groups for one K lane. + constexpr int64_t n_group_size = 8; + constexpr int64_t vnni_size = 4; + static_assert(N % (2 * n_group_size) == 0); + + int64_t n_group = n / n_group_size; + int64_t ni = n % n_group_size; + int64_t ki = k % vnni_size; + int64_t k_base = k - ki; + int64_t packed_n = (n_group / 2) * n_group_size + ni; + uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki]; + return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f); +} + +#if defined(CPU_CAPABILITY_RVV) +template +inline fixed_i8x8_t load_uint4_as_int8_rvv(const uint8_t* __restrict__ B, int64_t k) { + constexpr int64_t n_group_size = 8; + constexpr int64_t vnni_size = 4; + static_assert(N == 32); + static_assert(ldb == N / 2); + static_assert(group >= 0 && group < N / n_group_size); + + // Unpack: gather 8 packed int4 values from the VNNI4 layout. + const int64_t ki = k % vnni_size; + const int64_t k_base = k - ki; + constexpr int64_t packed_group = group / 2; + const uint8_t* packed_ptr = B + k_base * ldb + packed_group * n_group_size * vnni_size + ki; + + fixed_u8x8_t packed = RVVI(__riscv_vlse8_v_u8, LMUL_64)(packed_ptr, vnni_size, n_group_size); + if constexpr (group % 2 == 1) { + packed = RVVI(__riscv_vsrl_vx_u8, LMUL_64)(packed, 4, n_group_size); + } + fixed_u8x8_t nibbles = RVVI(__riscv_vand_vx_u8, LMUL_64)(packed, 0x0f, n_group_size); + return RVVI4(__riscv_vreinterpret_v_u8, LMUL_64, _i8, LMUL_64)(nibbles); +} + +inline fixed_i32x8_t gemm_accum_uint8_int8_rvv(fixed_i32x8_t acc, uint8_t a, fixed_i8x8_t b) { + constexpr int64_t vl = 8; + fixed_i16x8_t b_i16 = RVVI(__riscv_vsext_vf2_i16, LMUL_128)(b, vl); + return RVVI(__riscv_vwmacc_vx_i32, LMUL_256)(acc, static_cast(a), b_i16, vl); +} + +template +inline fixed_i32x8_t gemm_accum_uint4_rvv( + fixed_i32x8_t acc, + const uint8_t* __restrict__ B, + const int8_t* __restrict__ qzeros_b, + uint8_t a, + int64_t k) { + constexpr int64_t n_group_size = 8; + fixed_i8x8_t b = load_uint4_as_int8_rvv(B, k); + fixed_i8x8_t qzeros = + RVVI(__riscv_vle8_v_i8, LMUL_64)(qzeros_b + group * n_group_size, n_group_size); + b = RVVI(__riscv_vsub_vv_i8, LMUL_64)(b, qzeros, n_group_size); + return gemm_accum_uint8_int8_rvv(acc, a, b); +} + +template +inline void _dequant_and_store_rvv( + float* __restrict__ C, + fixed_i32x8_t acc, + const float* __restrict__ scales_a, + const int32_t* __restrict__ qzeros_a, + const float* __restrict__ scales_b, + const int32_t* __restrict__ compensation, + int64_t m, + int64_t ldc) { + constexpr int64_t n_group_size = 8; + constexpr int64_t n = group * n_group_size; + constexpr int64_t vl = n_group_size; + + // Dequant compensation: remove activation zero-point contribution. + fixed_i32x8_t comp = RVVI(__riscv_vle32_v_i32, LMUL_256)(compensation + n, vl); + fixed_i32x8_t zp_comp = RVVI(__riscv_vmul_vx_i32, LMUL_256)(comp, qzeros_a[m], vl); + acc = RVVI(__riscv_vsub_vv_i32, LMUL_256)(acc, zp_comp, vl); + + // Scale: convert int32 accumulators to fp32 and apply activation/weight scales. + fixed_fp32x8_t acc_f = RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(acc, vl); + acc_f = RVVI(__riscv_vfmul_vf_f32, LMUL_256)(acc_f, scales_a[m], vl); + fixed_fp32x8_t scale_b = RVVI(__riscv_vle32_v_f32, LMUL_256)(scales_b + n, vl); + acc_f = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(acc_f, scale_b, vl); + + // Store: accumulate into the float scratch buffer that already holds bias/zero. + float* c_ptr = C + m * ldc + n; + fixed_fp32x8_t c_old = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr, vl); + fixed_fp32x8_t c_new = RVVI(__riscv_vfadd_vv_f32, LMUL_256)(c_old, acc_f, vl); + RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr, c_new, vl); +} + +template +void _dequant_gemm_accum_rvv( + float* __restrict__ C, + const uint8_t* __restrict__ A, + const float* __restrict__ scales_a, + const int32_t* __restrict__ qzeros_a, + const uint8_t* __restrict__ B, + const float* __restrict__ scales_b, + const int8_t* __restrict__ qzeros_b, + const int32_t* __restrict__ compensation, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc) { + static_assert(N == 32); + static_assert(ldb == N / 2); + constexpr int64_t vl = 8; + + // Accumulate one C row over the 32-column block. + for (int64_t m = 0; m < M; ++m) { + fixed_i32x8_t acc0 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc1 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc2 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc3 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + // A[m][k] @ B[k][0:32] -> acc[m][0:32] + for (int64_t k = 0; k < K; ++k) { + // GEMM K step: one scalar activation updates four 8-column RVV tiles. + const uint8_t a = A[m * lda + k]; + acc0 = gemm_accum_uint4_rvv(acc0, B, qzeros_b, a, k); + acc1 = gemm_accum_uint4_rvv(acc1, B, qzeros_b, a, k); + acc2 = gemm_accum_uint4_rvv(acc2, B, qzeros_b, a, k); + acc3 = gemm_accum_uint4_rvv(acc3, B, qzeros_b, a, k); + } + + // Dequant/scale/store each 8-column group back into C. + _dequant_and_store_rvv<0>(C, acc0, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<1>(C, acc1, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<2>(C, acc2, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<3>(C, acc3, scales_a, qzeros_a, scales_b, compensation, m, ldc); + } +} +#endif + template void _dequant_gemm_accum( float* C, @@ -319,9 +455,31 @@ void _dequant_gemm_accum( _dequant_and_store( C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/); } else +#elif defined(CPU_CAPABILITY_RVV) + if constexpr (!sym_quant_act && N == BLOCK_N && ldb == BLOCK_N / 2) { + _dequant_gemm_accum_rvv(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, M, K, lda, ldc); + return; + } else #endif { - TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!"); + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + int32_t acc = 0; + for (int64_t k = 0; k < K; ++k) { + int32_t b = load_uint4_vnni(B, k, n) - qzeros_b[n]; + if constexpr (sym_quant_act) { + const int8_t* A_s8 = reinterpret_cast(A); + acc += static_cast(A_s8[m * lda + k]) * b; + } else { + acc += static_cast(A[m * lda + k]) * b; + } + } + if constexpr (!sym_quant_act) { + acc -= qzeros_a[m] * compensation[n]; + } + C[m * ldc + n] += static_cast(acc) * scales_a[m] * scales_b[n]; + } + } } } @@ -496,9 +654,11 @@ void _da8w4_linear_impl( store_out(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/); } } +#if defined(CPU_CAPABILITY_AVX512) if (use_brgemm) { at::native::cpublas::brgemm_release(); } +#endif }); } diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 77ffeec9fe7..407cfe60434 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -9,11 +9,19 @@ #define CPU_CAPABILITY_AVX512 #endif +#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) +#define CPU_CAPABILITY_RVV +#endif + #include #include #if defined(CPU_CAPABILITY_AVX512) #include #endif + +#if defined(CPU_CAPABILITY_RVV) +#include "../cpu_types_riscv_defs.hpp" +#endif namespace { using namespace at::vec; @@ -245,7 +253,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric for (int64_t k = 0; k < K; ++k) { const float val = static_cast(A[k]) * inv_scale; - Aq[k] = (uint8_t)(std::round(val)) + 128; + Aq[k] = static_cast(static_cast(std::round(val)) + 128); } As = scale; } diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp index a76b8bc6937..30192196b95 100644 --- a/csrc/cpu/spec_decode_utils.cpp +++ b/csrc/cpu/spec_decode_utils.cpp @@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl( } } +void copy_and_expand_dflash_inputs_kernel_impl( + const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, + torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, + torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, + torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, + const torch::Tensor& block_table, const torch::Tensor& query_start_loc, + const std::optional& num_rejected_tokens, + const int64_t parallel_drafting_token_id, const int64_t block_size, + const int64_t num_query_per_req, const int64_t num_speculative_tokens, + const int64_t total_input_tokens, const bool has_num_rejected) { + const int64_t num_reqs = query_start_loc.size(0) - 1; + + const int64_t* next_ids_ptr = next_token_ids.data_ptr(); + const int64_t* target_pos_ptr = target_positions.data_ptr(); + const int32_t* block_table_ptr = block_table.data_ptr(); + const int32_t* query_start_ptr = query_start_loc.data_ptr(); + const int64_t* rejected_ptr = + has_num_rejected && num_rejected_tokens.has_value() + ? num_rejected_tokens.value().data_ptr() + : nullptr; + + int64_t* out_ids_ptr = out_input_ids.data_ptr(); + int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr(); + int64_t* out_query_pos_ptr = out_query_positions.data_ptr(); + int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr(); + int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr(); + int32_t* out_token_idx_ptr = out_token_indices.data_ptr(); + + const int64_t block_table_stride = block_table.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int32_t ctx_start = query_start_ptr[req_idx]; + int32_t ctx_end = query_start_ptr[req_idx + 1]; + int64_t num_ctx = ctx_end - ctx_start; + int64_t valid_ctx_end = ctx_end; + if (rejected_ptr != nullptr) { + valid_ctx_end -= rejected_ptr[req_idx]; + } + // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that + // valid_ctx_end - 1 never reads before the request's context range. + valid_ctx_end = + std::max(valid_ctx_end, static_cast(ctx_start + 1)); + + int64_t last_pos = target_pos_ptr[valid_ctx_end - 1]; + + for (int64_t j = 0; j < num_ctx; ++j) { + int64_t ctx_idx = ctx_start + j; + int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1); + int64_t position = target_pos_ptr[ctx_pos_idx]; + int64_t block_num = position / block_size; + block_num = std::min(block_num, block_table_stride - 1); + int32_t block_id = + block_table_ptr[req_idx * block_table_stride + block_num]; + int64_t slot = block_id * block_size + (position % block_size); + + out_ctx_pos_ptr[ctx_idx] = position; + out_ctx_slot_ptr[ctx_idx] = slot; + } + + for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) { + int64_t query_out = req_idx * num_query_per_req + query_off; + int64_t position = last_pos + 1 + query_off; + int64_t block_num = position / block_size; + block_num = std::min(block_num, block_table_stride - 1); + int32_t block_id = + block_table_ptr[req_idx * block_table_stride + block_num]; + int64_t slot = block_id * block_size + (position % block_size); + + out_query_pos_ptr[query_out] = position; + out_query_slot_ptr[query_out] = slot; + out_ids_ptr[query_out] = + query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id; + + if (query_off > 0) { + int64_t sample_out_idx = + req_idx * num_speculative_tokens + (query_off - 1); + out_token_idx_ptr[sample_out_idx] = query_out; + } + } + } +} + 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, diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 7a8188b8c8c..50ce23d89b7 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -120,6 +120,14 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0, double softplus_threshold = 20.0); +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( + const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q, + const at::Tensor& k, const at::Tensor& v, const at::Tensor& a, + const at::Tensor& b, at::Tensor& initial_state_source, + const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens, + const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel, + double softplus_beta = 1.0, double softplus_threshold = 20.0); + std::tuple fused_gdn_gating_cpu( const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias); @@ -146,13 +154,16 @@ at::Tensor causal_conv1d_update_cpu( void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, const std::string& activation); +bool cpu_attn_has_isa(const std::string& isa); + 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, const torch::Tensor& seq_lens, at::ScalarType dtype, const torch::Tensor& query_start_loc, const bool casual, const int64_t window_size, const std::string& isa_hint, - const bool enable_kv_split); + const bool enable_kv_split, + const std::optional& dynamic_causal); void cpu_attn_reshape_and_cache(const torch::Tensor& key, const torch::Tensor& value, @@ -169,10 +180,10 @@ void cpu_attention_with_kv_cache( const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens, const double scale, const bool causal, const std::optional& alibi_slopes, - const int64_t sliding_window_left, const int64_t sliding_window_right, - const torch::Tensor& block_table, const double softcap, - const torch::Tensor& scheduler_metadata, - const std::optional& s_aux, const double k_scale, + const int64_t sliding_window_left, const torch::Tensor& block_table, + const double softcap, const torch::Tensor& scheduler_metadata, + const std::optional& s_aux, + const std::optional& dynamic_causal, const double k_scale, const double v_scale, const std::string& kv_cache_dtype); // Note: just for avoiding importing errors @@ -234,6 +245,16 @@ void copy_and_expand_eagle_inputs_kernel_impl( 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 copy_and_expand_dflash_inputs_kernel_impl( + const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, + torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, + torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, + torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, + const torch::Tensor& block_table, const torch::Tensor& query_start_loc, + const std::optional& num_rejected_tokens, + const int64_t parallel_drafting_token_id, const int64_t block_size, + const int64_t num_query_per_req, const int64_t num_speculative_tokens, + const int64_t total_input_tokens, const bool has_num_rejected); 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, @@ -265,7 +286,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "dynamic_4bit_int_moe(" "Tensor x, Tensor topk_ids, Tensor topk_weights," - "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2," + "Tensor w13_packed, Tensor w2_packed," + "int hidden_size, int intermediate_size," "int group_size, bool apply_router_weight_on_input, int activation_kind" ") -> Tensor"); @@ -285,6 +307,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul); + // GELU tanh implementation. + ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()"); + ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh); + // GELU implementation used in GPT-2. ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_new", torch::kCPU, &gelu_new); @@ -309,13 +335,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Layernorm // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( - "rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> " + "rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> " "()"); ops.impl("rms_norm", torch::kCPU, &rms_norm); // In-place fused Add and RMS Normalization. ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, " + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " "float epsilon) -> ()"); ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm); @@ -329,8 +355,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding); // Quantization -#if defined(__AVX512F__) || defined(__AVX2__) || \ - (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) +#if defined(__AVX512F__) || defined(__AVX2__) || \ + (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \ + defined(__riscv_v) // Helper function to release oneDNN handlers ops.def("release_dnnl_matmul_handler(int handler) -> ()", &release_dnnl_matmul_handler); @@ -428,19 +455,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("int8_scaled_mm_with_quant", torch::kCPU, &int8_scaled_mm_with_quant); - // Adapted from sglang: INT4 W4A8 kernels - ops.def( - "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor " - "scales, int quant_method_4bit) -> (Tensor, " - "Tensor, Tensor)"); - ops.impl("convert_weight_packed_scale_zp", torch::kCPU, - &convert_weight_packed_scale_zp); - - ops.def( - "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " - "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); - ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); - // Adapted from sglang: FP8 W8A16 kernel ops.def( "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) " @@ -467,6 +481,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif +#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ + defined(__AVX512VNNI__)) || \ + defined(__riscv) + // Adapted from sglang: INT4 W4A8 kernels + ops.def( + "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor " + "scales, int quant_method_4bit) -> (Tensor, " + "Tensor, Tensor)"); + ops.impl("convert_weight_packed_scale_zp", torch::kCPU, + &convert_weight_packed_scale_zp); + + ops.def( + "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " + "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); + ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); +#endif + // Adapted from sglang: GDN kernels ops.def( "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, " @@ -485,17 +516,27 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "softplus_threshold=20.0) -> Tensor"); ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU, &fused_sigmoid_gating_delta_rule_update_cpu); + ops.def( + "fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor " + "dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, " + "Tensor(a!) initial_state_source, Tensor spec_state_indices, " + "Tensor num_accepted_tokens, Tensor cu_seqlens, bool " + "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float " + "softplus_threshold=20.0) -> Tensor"); + ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU, + &fused_sigmoid_gating_delta_rule_update_spec_cpu); ops.def( "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) " "-> (Tensor, Tensor)"); ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); // CPU attention kernels + ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa); ops.def( "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " "query_start_loc, bool casual, int window_size, str isa_hint, bool " - "enable_kv_split) -> Tensor", + "enable_kv_split, Tensor? dynamic_causal) -> Tensor", &get_scheduler_metadata); ops.def( "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) " @@ -507,8 +548,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor " "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor " "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt " - "sliding_window_left, SymInt sliding_window_right, Tensor block_table, " - "float softcap, Tensor scheduler_metadata, Tensor? s_aux, " + "sliding_window_size, Tensor block_table, " + "float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? " + "dynamic_causal, " "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> " "()", &cpu_attention_with_kv_cache); @@ -528,7 +570,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif // fused moe -#if defined(__AVX512F__) +#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT)) ops.def( "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) " "-> ()"); @@ -539,7 +581,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#endif +#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT)) ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," @@ -589,6 +631,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "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( + "copy_and_expand_dflash_inputs_kernel_impl(" + "Tensor next_token_ids, Tensor target_positions, " + "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, " + "Tensor(a4!) out_query_positions, " + "Tensor(a5!) out_context_slot_mapping, " + "Tensor(a6!) out_query_slot_mapping, " + "Tensor(a7!) out_token_indices, Tensor block_table, " + "Tensor query_start_loc, Tensor? num_rejected_tokens, " + "SymInt parallel_drafting_token_id, SymInt block_size, " + "SymInt num_query_per_req, SymInt num_speculative_tokens, " + "SymInt total_input_tokens, bool has_num_rejected) -> ()", + &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl); ops.def( "rejection_greedy_sample_kernel_impl(" "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp index 394e67e3a03..78ee7081b24 100644 --- a/csrc/cpu/utils.hpp +++ b/csrc/cpu/utils.hpp @@ -2,19 +2,24 @@ #define UTILS_HPP #include +#include #include #include #include "cpu/cpu_types.hpp" namespace cpu_utils { -enum class ISA { AMX, VEC }; +enum class ISA { AMX, VEC, RVV, NEON }; inline ISA get_isa(const std::string& isa) { if (isa == "amx") { return ISA::AMX; } else if (isa == "vec") { return ISA::VEC; + } else if (isa == "rvv") { + return ISA::RVV; + } else if (isa == "neon") { + return ISA::NEON; } else { TORCH_CHECK(false, "Invalid isa type: " + isa); } @@ -71,14 +76,14 @@ inline int64_t get_available_l2_size() { if (l2_cache_size == 0) { l2_cache_size = 256 * 1024; } - return static_cast(l2_cache_size) >> 1; // use 50% of L2 cache + return static_cast(l2_cache_size) >> 1; }(); 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 l2_cache_size >> 1; }(); return size; #endif diff --git a/csrc/cuda_view.cu b/csrc/cuda_view.cu deleted file mode 100644 index 00e06a9329f..00000000000 --- a/csrc/cuda_view.cu +++ /dev/null @@ -1,60 +0,0 @@ -// TODO: Remove this once ROCm upgrade to torch 2.11. -#include -#include -#include - -// This function assumes that `cpu_tensor` is a CPU tensor, -// and that UVA (Unified Virtual Addressing) is enabled. -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) { - TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU"); - - // handle empty tensor - if (cpu_tensor.numel() == 0) { - return torch::empty(cpu_tensor.sizes(), - cpu_tensor.options().device(torch::kCUDA)); - } - - if (cpu_tensor.is_pinned()) { - // If CPU tensor is pinned, directly get the device pointer. - void* host_ptr = const_cast(cpu_tensor.data_ptr()); - void* device_ptr = nullptr; - cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); - TORCH_CHECK(err == cudaSuccess, - "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err)); - - return torch::from_blob( - device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), - [base = cpu_tensor](void*) {}, // keep cpu tensor alive - cpu_tensor.options().device(torch::kCUDA)); - } - - // If CPU tensor is not pinned, allocate a new pinned memory buffer. - torch::Tensor contiguous_cpu = cpu_tensor.contiguous(); - size_t nbytes = contiguous_cpu.nbytes(); - - void* host_ptr = nullptr; - cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped); - if (err != cudaSuccess) { - AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err)); - } - - err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes, - cudaMemcpyDefault); - if (err != cudaSuccess) { - cudaFreeHost(host_ptr); - AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err)); - } - - void* device_ptr = nullptr; - err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); - if (err != cudaSuccess) { - cudaFreeHost(host_ptr); - AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err)); - } - - auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); }; - - return torch::from_blob(device_ptr, contiguous_cpu.sizes(), - contiguous_cpu.strides(), deleter, - contiguous_cpu.options().device(torch::kCUDA)); -} \ No newline at end of file diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 73333f7125f..2329d51a149 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -48,8 +48,8 @@ static inline unsigned long long my_min(unsigned long long a, } static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size, - size_t alignment) { - CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0); + size_t alignment, CUdeviceptr addr = 0) { + CUresult status = cuMemAddressReserve(d_mem, size, alignment, addr, 0); if (status == CUresult(0) || alignment == 0) { return status; } @@ -58,7 +58,7 @@ static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size, // alignment even when physical VRAM is free. Let HIP choose the default // alignment, then verify that the returned address still satisfies the // requested alignment before accepting it. - status = cuMemAddressReserve(d_mem, size, 0, 0, 0); + status = cuMemAddressReserve(d_mem, size, 0, addr, 0); if (status != CUresult(0)) { return status; } @@ -535,7 +535,14 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) { Py_DECREF(py_result); PyGILState_Release(gstate); - unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, num_chunks); + // An empty chunk list means this allocation is asleep: its physical chunks + // were already unmapped and released by sleep(), but the virtual address is + // still held as a placeholder reservation. Skip unmap/release (freeing the + // placeholder address happens below). + if (num_chunks > 0) { + unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, + num_chunks); + } #else // Non-ROCm path: simple integer handle already extracted; drop temporary // Python refs while still holding the GIL, then release it. @@ -548,11 +555,13 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) { unmap_and_release(device, size, d_mem, p_memHandle); #endif - // free address and the handle + // Free the virtual address. On ROCm this also covers an asleep allocation, + // whose placeholder reservation made by sleep() is still held here. CUDA_CHECK(cuMemAddressFree(d_mem, size)); #ifndef USE_ROCM free(p_memHandle); #else + // Only awake allocations have per-chunk handles to free. for (auto i = 0; i < num_chunks; ++i) { free(p_memHandle[i]); } @@ -672,6 +681,29 @@ static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) { unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes, num_chunks); + // On ROCm/Linux, physical VRAM is only reclaimed once the virtual address + // range is freed; hipMemUnmap + hipMemRelease alone leave the memory + // resident (see ROCm#6021). Free the address to release physical memory, + // then immediately re-reserve the SAME address as an empty placeholder so + // the regular allocator cannot hand it out while we sleep. wake_up remaps + // physical chunks into this placeholder. + if (error_code == no_error) { + CUDA_CHECK(cuMemAddressFree(d_mem_ptr, recv_size)); + if (error_code == no_error) { + CUdeviceptr reserved = 0; + CUDA_CHECK(reserve_rocm_address(&reserved, recv_size, /*alignment=*/0, + d_mem_ptr)); + if (error_code == no_error && reserved != d_mem_ptr) { + (void)cuMemAddressFree(reserved, recv_size); + snprintf(error_msg, sizeof(error_msg), + "failed to re-reserve placeholder address on sleep " + "(requested %#llx, got %#llx)", + (unsigned long long)d_mem_ptr, (unsigned long long)reserved); + error_code = CUresult(1); + } + } + } + free(p_memHandle); free(chunk_sizes); #endif @@ -736,6 +768,7 @@ static PyObject* python_create_and_map(PyObject* self, PyObject* args) { chunk_sizes[i] = PyLong_AsUnsignedLongLong(size_py); } + // Address already reserved as a placeholder by sleep(); just remap chunks. create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes, num_chunks); diff --git a/csrc/custom_all_reduce_test.cu b/csrc/custom_all_reduce_test.cu deleted file mode 100644 index f7f0823465d..00000000000 --- a/csrc/custom_all_reduce_test.cu +++ /dev/null @@ -1,361 +0,0 @@ -/** - * This is a standalone test for custom allreduce. - * To compile, make sure you have MPI and NCCL installed in your system. - * export MPI_HOME=XXX - * nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o - * custom_all_reduce_test -lnccl -I${MPI_HOME}/include -lmpi - * - * Warning: this C++ test is not designed to be very readable and was used - * during the rapid prototyping process. - * - * To run: - * mpirun --allow-run-as-root -np 8 ./custom_all_reduce_test - */ -#include -#include -#include -#include - -#include -#include - -#include "cuda_profiler_api.h" -#include "custom_all_reduce.cuh" -#include "mpi.h" -#ifdef USE_ROCM - #include -typedef __hip_bfloat16 nv_bfloat16; - #include "rccl/rccl.h" - #include "custom_all_reduce_hip.cuh" -#else - #include "nccl.h" - #include "custom_all_reduce.cuh" -#endif - -#define MPICHECK(cmd) \ - do { \ - int e = cmd; \ - if (e != MPI_SUCCESS) { \ - printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -#define NCCLCHECK(cmd) \ - do { \ - ncclResult_t r = cmd; \ - if (r != ncclSuccess) { \ - printf("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, \ - ncclGetErrorString(r)); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -#ifdef USE_ROCM -__global__ void dummy_kernel() { - for (int i = 0; i < 100; i++) { - uint64_t start = wall_clock64(); - uint64_t cycles_elapsed; - do { - cycles_elapsed = wall_clock64() - start; - } while (cycles_elapsed < 100); - } - for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms -} -#else -__global__ void dummy_kernel() { - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms - #else - for (int i = 0; i < 100; i++) { - long long int start = clock64(); - while (clock64() - start < 150000000); // approximately 98.4ms on P40 - } - #endif -} -#endif - -template -__global__ void set_data(T* data, int size, int myRank) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - data[idx] = myRank * 0.11f; - } -} - -template -__global__ void convert_data(const T* data1, const T* data2, double* fdata1, - double* fdata2, int size) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - fdata1[idx] = data1[idx]; - fdata2[idx] = data2[idx]; - } -} - -__global__ void init_rand(curandState_t* state, int size, int nRanks) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - for (int i = 0; i < nRanks; i++) { - curand_init(i + 1, idx, 0, &state[idx * nRanks + i]); - } - } -} - -template -__global__ void gen_data(curandState_t* state, T* data, double* ground_truth, - int myRank, int nRanks, int size) { - for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; - idx += gridDim.x * blockDim.x) { - double sum = 0.0; - for (int i = 0; i < nRanks; i++) { - double val = curand_uniform_double(&state[idx * nRanks + i]) * 4; - T hval = val; // downcast first - sum += static_cast(hval); - if (i == myRank) data[idx] = hval; - } - ground_truth[idx] = sum; - } -} - -template -void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit, - int data_size, bool performance_test) { - T* result; - cudaStream_t stream; - CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); - CUDACHECK(cudaMalloc(&result, data_size * sizeof(T))); - CUDACHECK(cudaMemset(result, 0, data_size * sizeof(T))); - - cudaIpcMemHandle_t self_data_handle; - cudaIpcMemHandle_t data_handles[8]; - vllm::Signal* buffer; - T* self_data_copy; - /** - * Allocate IPC buffer - * - * The first section is a temporary buffer for storing intermediate allreduce - * results, if a particular algorithm requires it. The second section is for - * the input to the allreduce. The actual API takes the input pointer as an - * argument (that is, they can and usually should be allocated separately). - * But since the input pointers and the temporary buffer all require IPC - * registration, they are allocated and registered together in the test for - * convenience. - */ -#ifdef USE_ROCM - CUDACHECK(hipExtMallocWithFlags( - (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal), - hipDeviceMallocUncached)); -#else - CUDACHECK( - cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal))); -#endif - CUDACHECK( - cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal))); - CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T))); - CUDACHECK(cudaIpcGetMemHandle(&self_data_handle, buffer)); - - MPICHECK(MPI_Allgather(&self_data_handle, sizeof(cudaIpcMemHandle_t), - MPI_BYTE, data_handles, sizeof(cudaIpcMemHandle_t), - MPI_BYTE, MPI_COMM_WORLD)); - - void* rank_data; - size_t rank_data_sz = 16 * 1024 * 1024; - CUDACHECK(cudaMalloc(&rank_data, rank_data_sz)); - vllm::Signal* ipc_ptrs[8]; - for (int i = 0; i < nRanks; i++) { - if (i == myRank) - ipc_ptrs[i] = buffer; - else - CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i], - cudaIpcMemLazyEnablePeerAccess)); - } - vllm::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks); - auto* self_data = - reinterpret_cast(reinterpret_cast(buffer) + - sizeof(vllm::Signal) + data_size * sizeof(T)); - // hack buffer registration - { - void* data[8]; - for (int i = 0; i < nRanks; i++) { - data[i] = - ((char*)ipc_ptrs[i]) + sizeof(vllm::Signal) + data_size * sizeof(T); - } - fa.register_buffer(data); - } - - double* ground_truth; - CUDACHECK(cudaMallocHost(&ground_truth, data_size * sizeof(double))); - curandState_t* states; - CUDACHECK(cudaMalloc(&states, sizeof(curandState_t) * nRanks * data_size)); - init_rand<<<108, 1024, 0, stream>>>(states, data_size, nRanks); - gen_data<<<108, 1024, 0, stream>>>(states, self_data, ground_truth, myRank, - nRanks, data_size); - CUDACHECK(cudaMemcpyAsync(self_data_copy, self_data, data_size * sizeof(T), - cudaMemcpyDeviceToDevice, stream)); - cudaEvent_t start, stop; - CUDACHECK(cudaEventCreate(&start)); - CUDACHECK(cudaEventCreate(&stop)); - - ncclDataType_t ncclDtype; - if (std::is_same::value) { - ncclDtype = ncclFloat16; - } else if (std::is_same::value) { - ncclDtype = ncclBfloat16; - } else { - ncclDtype = ncclFloat; - } - double *nccl_result, *my_result; - CUDACHECK(cudaMallocHost(&nccl_result, data_size * sizeof(double))); - CUDACHECK(cudaMallocHost(&my_result, data_size * sizeof(double))); - if (performance_test) { - dummy_kernel<<<1, 1, 0, stream>>>(); - constexpr int warmup_iters = 5; - constexpr int num_iters = 100; - // warmup - for (int i = 0; i < warmup_iters; i++) { - NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum, - comm, stream)); - } - CUDACHECK(cudaEventRecord(start, stream)); - for (int i = 0; i < num_iters; i++) { - NCCLCHECK(ncclAllReduce(result, result, data_size, ncclDtype, ncclSum, - comm, stream)); - } - CUDACHECK(cudaEventRecord(stop, stream)); - CUDACHECK(cudaStreamSynchronize(stream)); - float allreduce_ms = 0; - cudaEventElapsedTime(&allreduce_ms, start, stop); - - dummy_kernel<<<1, 1, 0, stream>>>(); - // warm up - for (int i = 0; i < warmup_iters; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - } - CUDACHECK(cudaEventRecord(start, stream)); - for (int i = 0; i < num_iters; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - } - CUDACHECK(cudaEventRecord(stop, stream)); - CUDACHECK(cudaStreamSynchronize(stream)); - - float duration_ms = 0; - cudaEventElapsedTime(&duration_ms, start, stop); - if (myRank == 0) - printf( - "Rank %d done, nGPUs:%d, sz (kb): %d, %d, %d, my time:%.2fus, nccl " - "time:%.2fus\n", - myRank, nRanks, data_size * sizeof(T) / 1024, threads, block_limit, - duration_ms * 1e3 / num_iters, allreduce_ms * 1e3 / num_iters); - - // And wait for all the queued up work to complete - CUDACHECK(cudaStreamSynchronize(stream)); - - NCCLCHECK(ncclAllReduce(self_data_copy, self_data, data_size, ncclDtype, - ncclSum, comm, stream)); - - convert_data<<<108, 1024, 0, stream>>>(self_data, result, nccl_result, - my_result, data_size); - CUDACHECK(cudaStreamSynchronize(stream)); - - for (unsigned long j = 0; j < data_size; j++) { - auto diff = abs(nccl_result[j] - my_result[j]); - if (diff >= 4e-2) { - printf("Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n", - myRank, j, nccl_result[j], my_result[j], ground_truth[j]); - break; - } - } - long double nccl_diffs = 0.0; - long double my_diffs = 0.0; - for (int j = 0; j < data_size; j++) { - nccl_diffs += abs(nccl_result[j] - ground_truth[j]); - my_diffs += abs(my_result[j] - ground_truth[j]); - } - if (myRank == 0) - std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size - << " me: " << my_diffs / data_size << std::endl; - } else { - for (int i = 0; i < 100; i++) { - fa.allreduce(stream, self_data, result, data_size, threads, - block_limit); - CUDACHECK(cudaStreamSynchronize(stream)); - NCCLCHECK(ncclAllReduce(self_data, self_data_copy, data_size, ncclDtype, - ncclSum, comm, stream)); - convert_data<<<108, 1024, 0, stream>>>( - self_data_copy, result, nccl_result, my_result, data_size); - CUDACHECK(cudaStreamSynchronize(stream)); - - for (unsigned long j = 0; j < data_size; j++) { - auto diff = abs(nccl_result[j] - my_result[j]); - if (diff >= 4e-2) { - printf( - "Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n", - myRank, j, nccl_result[j], my_result[j], ground_truth[j]); - break; - } - } - } - if (myRank == 0) - printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks, - data_size * sizeof(T) / 1024, threads, block_limit); - // long double nccl_diffs = 0.0; - // long double my_diffs = 0.0; - // for (int j = 0; j < data_size; j++) { - // nccl_diffs += abs(nccl_result[j] - ground_truth[j]); - // my_diffs += abs(my_result[j] - ground_truth[j]); - // } - // if (myRank == 0) - // std::cout << "average abs diffs: nccl: " << nccl_diffs / data_size - // << " me: " << my_diffs / data_size << std::endl; - } - - CUDACHECK(cudaFree(result)); - CUDACHECK(cudaFree(self_data_copy)); - CUDACHECK(cudaFree(rank_data)); - CUDACHECK(cudaFree(buffer)); - CUDACHECK(cudaFree(states)); - CUDACHECK(cudaFreeHost(ground_truth)); - CUDACHECK(cudaFreeHost(nccl_result)); - CUDACHECK(cudaFreeHost(my_result)); - CUDACHECK(cudaStreamDestroy(stream)); -} - -int main(int argc, char** argv) { - int nRanks, myRank; - MPICHECK(MPI_Init(&argc, &argv)); - MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank)); - MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks)); - CUDACHECK(cudaSetDevice(myRank)); - ncclUniqueId id; - ncclComm_t comm; - if (myRank == 0) ncclGetUniqueId(&id); - MPICHECK(MPI_Bcast(static_cast(&id), sizeof(id), MPI_BYTE, 0, - MPI_COMM_WORLD)); - NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank)); - - bool performance_test = true; - cudaProfilerStart(); -// Uncomment to scan through different block size configs. -// for (int threads : {256, 512, 1024}) { -// for (int block_limit = 16; block_limit < 112; block_limit += 4) { -// run(myRank, nRanks, comm, threads, block_limit, 1024 * 1024, -// performance_test); -// } -// } -#ifdef USE_ROCM - const int block_limit = 16; -#else - const int block_limit = 36; -#endif - // Scan through different sizes to test performance. - for (int sz = 512; sz <= (8 << 20); sz *= 2) { - run(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test); - } - - cudaProfilerStop(); - MPICHECK(MPI_Finalize()); - return EXIT_SUCCESS; -} \ No newline at end of file diff --git a/csrc/custom_quickreduce.cu b/csrc/custom_quickreduce.cu index 33d0d4a7226..d4e5d179a54 100644 --- a/csrc/custom_quickreduce.cu +++ b/csrc/custom_quickreduce.cu @@ -97,18 +97,28 @@ int64_t qr_max_size() { cast_bf2half>; \ template struct quickreduce::AllReduceTwoshot, cast_bf2half>; + // INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the + // world_size == 2 kernel for it. + #define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \ + template struct quickreduce::AllReduceTwoshot, cast_bf2half>; + INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16, + quickreduce::CodecQ3, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16, + quickreduce::CodecQ3, true) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false) #endif // USE_ROCM \ No newline at end of file diff --git a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py index 34fb64c413d..d692502f3ff 100644 --- a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py +++ b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py @@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = { } VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = { - DataType.u8: "at::ScalarType::Byte", - DataType.s8: "at::ScalarType::Char", - DataType.e4m3: "at::ScalarType::Float8_e4m3fn", - DataType.s32: "at::ScalarType::Int", - DataType.f16: "at::ScalarType::Half", - DataType.bf16: "at::ScalarType::BFloat16", - DataType.f32: "at::ScalarType::Float", + DataType.u8: "torch::headeronly::ScalarType::Byte", + DataType.s8: "torch::headeronly::ScalarType::Char", + DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn", + DataType.s32: "torch::headeronly::ScalarType::Int", + DataType.f16: "torch::headeronly::ScalarType::Half", + DataType.bf16: "torch::headeronly::ScalarType::BFloat16", + DataType.f32: "torch::headeronly::ScalarType::Float", } VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = { diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp new file mode 100644 index 00000000000..fdf3e614e64 --- /dev/null +++ b/csrc/fs_io.cpp @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include + +#include + +#include + +extern "C" { + +static void _batch_lookup(const std::vector& paths, + std::vector& exists_flags) { + for (size_t i = 0; i < paths.size(); i++) { + exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; + } +} + +/// @brief Check file existence for a batch of paths. +/// @param paths list[str] – absolute paths to check. +/// @return list[bool] – True if the corresponding path exists, False otherwise. +/// @note Releases the GIL for the entire batch. File existence via access(2). +static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { + PyObject* path_list; + if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(path_list); + std::vector paths(n); + for (Py_ssize_t i = 0; i < n; i++) { + paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr); + if (paths[i] == nullptr) { + return nullptr; + } + } + + std::vector exists_flags(n); + { + Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags); + Py_END_ALLOW_THREADS + } + + PyObject* result = PyList_New(n); + if (result == nullptr) { + return nullptr; + } + for (Py_ssize_t i = 0; i < n; i++) { + PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i])); + } + return result; +} + +static PyMethodDef fs_io_C_methods[] = { + {"batch_lookup", batch_lookup, METH_VARARGS, + "batch_lookup(paths: list[str]) -> list[bool]\n" + "\n" + "Check file existence for a batch of paths."}, + {nullptr, nullptr, 0, nullptr}, +}; + +static struct PyModuleDef fs_io_C_module = { + PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1, + fs_io_C_methods, +}; + +PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); } + +} // extern "C" diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index cdab456348e..60b8ca5f382 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -10,11 +10,20 @@ namespace vllm { -template __device__ __forceinline__ scalar_t compute(const scalar_t& x, const scalar_t& y, - const float limit) { + const float limit, + const float alpha, + const float beta) { if constexpr (act_first) { scalar_t gate = x; scalar_t up = y; @@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x, gate = (scalar_t)fminf((float)gate, limit); up = (scalar_t)fmaxf(fminf((float)up, limit), -limit); } - return ACT_FN(gate) * up; + // act_first: gate is the activated half -> alpha applies to gate; + // beta is added to up (the non-activated half). + return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta)); } else { scalar_t gate = x; scalar_t up = y; @@ -30,55 +41,68 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x, gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit); up = (scalar_t)fminf((float)up, limit); } - return gate * ACT_FN(up); + // !act_first: up is the activated half -> alpha applies to up; + // beta is added to gate (the non-activated half). + return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha)); } } -template __device__ __forceinline__ packed_t packed_compute(const packed_t& x, const packed_t& y, - const float limit) { + const float limit, + const float alpha, + const float beta) { if constexpr (act_first) { packed_t gate = x; packed_t up = y; + float2 u = cast_to_float2(up); 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); + // act_first: gate is the activated half -> alpha applies to gate; + // beta is added to up (the non-activated half). + float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha)); + activated.x *= u.x + beta; + activated.y *= u.y + beta; + return cast_to_packed(activated); } else { packed_t gate = x; packed_t up = y; + float2 g = cast_to_float2(gate); 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)); + // !act_first: up is the activated half -> alpha applies to up; + // beta is added to gate (the non-activated half). + float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha)); + activated.x *= g.x + beta; + activated.y *= g.y + beta; + return cast_to_packed(activated); } } // Activation and gating kernel template. template + scalar_t (*ACT_FN)(const scalar_t&, const float), + packed_t (*PACKED_ACT_FN)(const packed_t&, const float), + bool act_first, 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 float limit) { + const int d, const float limit, const float alpha, const float beta) { 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; @@ -105,7 +129,7 @@ __global__ void act_and_mul_kernel( for (int j = 0; j < pvec_t::NUM_ELTS; j++) { x.elts[j] = packed_compute( - x.elts[j], y.elts[j], limit); + x.elts[j], y.elts[j], limit, alpha, beta); } if constexpr (use_256b) { st256(x, &out_vec[i]); @@ -118,29 +142,34 @@ __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, limit); + out_ptr[idx] = compute( + x, y, limit, alpha, beta); } } } +// Gated activations take an `alpha` argument that scales the sigmoid input +// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which +// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a +// non-default alpha. Activations that do not use alpha simply ignore it. template -__device__ __forceinline__ T silu_kernel(const T& x) { - // x * sigmoid(x) - return (T)(((float)x) / (1.0f + expf((float)-x))); +__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) { + // x * sigmoid(alpha * x) + return (T)(((float)x) / (1.0f + expf((float)-x * alpha))); } template -__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) { - // x * sigmoid(x) +__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val, + const float alpha) { + // x * sigmoid(alpha * x) float2 fval = cast_to_float2(val); - fval.x = fval.x / (1.0f + expf(-fval.x)); - fval.y = fval.y / (1.0f + expf(-fval.y)); + fval.x = fval.x / (1.0f + expf(-fval.x * alpha)); + fval.y = fval.y / (1.0f + expf(-fval.y * alpha)); return cast_to_packed(fval); } template -__device__ __forceinline__ T gelu_kernel(const T& x) { +__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) { // Equivalent to PyTorch GELU with 'none' approximation. // Refer to: // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 @@ -150,7 +179,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) { } template -__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) { +__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val, + const float /*alpha*/) { // Equivalent to PyTorch GELU with 'none' approximation. // Refer to: // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38 @@ -162,7 +192,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) { } template -__device__ __forceinline__ T gelu_tanh_kernel(const T& x) { +__device__ __forceinline__ T gelu_tanh_kernel(const T& x, + const float /*alpha*/) { // Equivalent to PyTorch GELU with 'tanh' approximation. // Refer to: // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 @@ -176,7 +207,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) { template __device__ __forceinline__ packed_t -packed_gelu_tanh_kernel(const packed_t& val) { +packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) { // Equivalent to PyTorch GELU with 'tanh' approximation. // Refer to: // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30 @@ -202,7 +233,7 @@ packed_gelu_tanh_kernel(const packed_t& val) { // 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) \ + HAS_CLAMP, LIMIT, ALPHA, BETA) \ auto dtype = input.scalar_type(); \ int d = input.size(-1) / 2; \ int64_t num_tokens = input.numel() / input.size(-1); \ @@ -230,7 +261,7 @@ packed_gelu_tanh_kernel(const packed_t& val) { PACKED_KERNEL::Type>, \ ACT_FIRST, true, HAS_CLAMP, true><<>>( \ out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT); \ + input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ }); \ } else { \ VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \ @@ -240,7 +271,7 @@ packed_gelu_tanh_kernel(const packed_t& val) { PACKED_KERNEL::Type>, \ ACT_FIRST, true, HAS_CLAMP, false><<>>( \ out.mutable_data_ptr(), \ - input.const_data_ptr(), d, LIMIT); \ + input.const_data_ptr(), d, LIMIT, ALPHA, BETA); \ }); \ } \ } else { \ @@ -252,7 +283,7 @@ packed_gelu_tanh_kernel(const packed_t& val) { PACKED_KERNEL::Type>, \ ACT_FIRST, false, HAS_CLAMP><<>>( \ out.mutable_data_ptr(), input.const_data_ptr(), \ - d, LIMIT); \ + d, LIMIT, ALPHA, BETA); \ }); \ } @@ -260,14 +291,18 @@ void silu_and_mul(torch::stable::Tensor& out, // [..., d] torch::stable::Tensor& input) // [..., 2 * d] { LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel, - true, false, 0.0f); + true, false, 0.0f, 1.0f, 0.0f); } void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d] torch::stable::Tensor& input, // [..., 2 * d] - double limit) { + double limit, double alpha, double beta) { + // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit))) + // * (up.clamp(+-limit) + beta) + // alpha=1.0, beta=0.0 reduce this to silu(gate) * up. LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel, - true, true, (float)limit); + true, true, (float)limit, (float)alpha, + (float)beta); } void mul_and_silu(torch::stable::Tensor& out, // [..., d] @@ -276,21 +311,22 @@ void mul_and_silu(torch::stable::Tensor& out, // [..., d] // The difference between mul_and_silu and silu_and_mul is that mul_and_silu // applies the silu to the latter half of the input. LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel, - false, false, 0.0f); + false, false, 0.0f, 1.0f, 0.0f); } void gelu_and_mul(torch::stable::Tensor& out, // [..., d] torch::stable::Tensor& input) // [..., 2 * d] { LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel, - true, false, 0.0f); + true, false, 0.0f, 1.0f, 0.0f); } void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d] torch::stable::Tensor& input) // [..., 2 * d] { - LAUNCH_ACTIVATION_GATE_KERNEL( - vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f); + LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, + vllm::packed_gelu_tanh_kernel, true, false, + 0.0f, 1.0f, 0.0f); } namespace vllm { @@ -633,6 +669,14 @@ __device__ __forceinline__ T gelu_quick_kernel(const T& x) { return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x))); } +template +__device__ __forceinline__ T relu_squared_kernel(const T& x) { + // relu(x)^2 — introduced in https://arxiv.org/abs/2109.08668v2 + const float f = (float)x; + const float val = f > 0.0f ? f : 0.0f; + return (T)(val * val); +} + } // namespace vllm void gelu_new(torch::stable::Tensor& out, // [..., d] @@ -652,3 +696,9 @@ void gelu_quick(torch::stable::Tensor& out, // [..., d] { LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel); } + +void relu_squared(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input) // [..., d] +{ + LAUNCH_ACTIVATION_KERNEL(vllm::relu_squared_kernel); +} diff --git a/csrc/libtorch_stable/attention/attention_kernels.cuh b/csrc/libtorch_stable/attention/attention_kernels.cuh deleted file mode 100644 index c5f9a9876c3..00000000000 --- a/csrc/libtorch_stable/attention/attention_kernels.cuh +++ /dev/null @@ -1,667 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, 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 "../../attention/attention_dtypes.h" -#include "attention_utils.cuh" -#include "../../cuda_compat.h" - -#ifdef USE_ROCM - #include - #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh" -typedef __hip_bfloat16 __nv_bfloat16; -#else - #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh" -#endif - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -namespace vllm { - -// Utility function for attention softmax. -template -inline __device__ float block_sum(float* red_smem, float sum) { - // Decompose the thread index into warp / lane. - int warp = threadIdx.x / WARP_SIZE; - int lane = threadIdx.x % WARP_SIZE; - - // Compute the sum per warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - sum += VLLM_SHFL_XOR_SYNC(sum, mask); - } - - // Warp leaders store the data to shared memory. - if (lane == 0) { - red_smem[warp] = sum; - } - - // Make sure the data is in shared memory. - __syncthreads(); - - // The warps compute the final sums. - if (lane < NUM_WARPS) { - sum = red_smem[lane]; - } - - // Parallel reduction inside the warp. -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - sum += VLLM_SHFL_XOR_SYNC(sum, mask); - } - - // Broadcast to other threads. - return VLLM_SHFL_SYNC(sum, 0); -} - -// TODO(woosuk): Merge the last two dimensions of the grid. -// Grid: (num_heads, num_seqs, max_num_partitions). -template // Zero means no partitioning. -__device__ void paged_attention_kernel( - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - const int seq_idx = blockIdx.y; - const int partition_idx = blockIdx.z; - const int max_num_partitions = gridDim.z; - constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0; - const int seq_len = seq_lens[seq_idx]; - if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) { - // No work to do. Terminate the thread block. - return; - } - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int num_blocks_per_partition = - USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks; - - // [start_block_idx, end_block_idx) is the range of blocks to process. - const int start_block_idx = - USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0; - const int end_block_idx = - MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks); - const int num_blocks = end_block_idx - start_block_idx; - - // [start_token_idx, end_token_idx) is the range of tokens to process. - const int start_token_idx = start_block_idx * BLOCK_SIZE; - const int end_token_idx = - MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len); - const int num_tokens = end_token_idx - start_token_idx; - - constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1); - constexpr int NUM_THREAD_GROUPS = - NUM_THREADS / THREAD_GROUP_SIZE; // Note: This assumes THREAD_GROUP_SIZE - // divides NUM_THREADS - assert(NUM_THREADS % THREAD_GROUP_SIZE == 0); - constexpr int NUM_TOKENS_PER_THREAD_GROUP = - DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE); - constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE; - const int thread_idx = threadIdx.x; - const int warp_idx = thread_idx / WARP_SIZE; - const int lane = thread_idx % WARP_SIZE; - - const int head_idx = blockIdx.x; - const int num_heads = gridDim.x; - const int num_queries_per_kv = num_heads / num_kv_heads; - const int kv_head_idx = head_idx / num_queries_per_kv; - const float alibi_slope = - alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx]; - - // A vector type to store a part of a key or a query. - // The vector size is configured in such a way that the threads in a thread - // group fetch or compute 16 bytes at a time. For example, if the size of a - // thread group is 4 and the data type is half, then the vector size is 16 / - // (4 * sizeof(half)) == 2. - constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1); - using K_vec = typename Vec::Type; - using Q_vec = typename Vec::Type; - using Quant_vec = typename Vec::Type; - - constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE; - constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE; - - const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE; - const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE; - - // Load the query to registers. - // Each thread in a thread group has a different part of the query. - // For example, if the thread group size is 4, then the first thread in - // the group has 0, 4, 8, ... th vectors of the query, and the second thread - // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because - // q is split from a qkv tensor, it may not be contiguous. - const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE; - __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD]; -#pragma unroll - for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD; - i += NUM_THREAD_GROUPS) { - const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE; - q_vecs[thread_group_offset][i] = - *reinterpret_cast(q_ptr + vec_idx * VEC_SIZE); - } - __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a - // memory wall right before we use q_vecs - - // Memory planning. - extern __shared__ char shared_mem[]; - // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy. - float* logits = reinterpret_cast(shared_mem); - // Workspace for reduction. - __shared__ float red_smem[2 * NUM_WARPS]; - - // x == THREAD_GROUP_SIZE * VEC_SIZE - // Each thread group fetches x elements from the key at a time. - constexpr int x = 16 / sizeof(cache_t); - float qk_max = -FLT_MAX; - - // Iterate over the key blocks. - // Each warp fetches a block of keys for each iteration. - // Each thread group in a warp fetches a key from the block, and computes - // dot product with the query. - const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; - - // blocksparse specific vars - int bs_block_offset; - int q_bs_block_id; - if constexpr (IS_BLOCK_SPARSE) { - // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len, - // blocksparse_block_size); - q_bs_block_id = (seq_len - 1) / blocksparse_block_size; - if (blocksparse_head_sliding_step >= 0) - // sliding on q heads - bs_block_offset = - (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1; - else - // sliding on kv heads - bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) * - (-blocksparse_head_sliding_step) + - 1; - } - - for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; - block_idx += NUM_WARPS) { - // NOTE(woosuk): The block number is stored in int32. However, we cast it to - // int64 because int32 can lead to overflow when this variable is multiplied - // by large numbers (e.g., kv_block_stride). - // For blocksparse attention: skip computation on blocks that are not - // attended - if constexpr (IS_BLOCK_SPARSE) { - const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; - const bool is_remote = - ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0); - const bool is_local = - (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks); - if (!is_remote && !is_local) { - for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { - const int physical_block_offset = - (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - - if (thread_group_offset == 0) { - // NOTE(linxihui): assign very large number to skipped tokens to - // avoid contribution to the sumexp softmax normalizer. This will - // not be used at computing sum(softmax*v) as the blocks will be - // skipped. - logits[token_idx - start_token_idx] = -FLT_MAX; - } - } - continue; - } - } - const int64_t physical_block_number = - static_cast(block_table[block_idx]); - - // Load a key to registers. - // Each thread in a thread group has a different part of the key. - // For example, if the thread group size is 4, then the first thread in - // the group has 0, 4, 8, ... th vectors of the key, and the second thread - // has 1, 5, 9, ... th vectors of the key, and so on. - for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { - const int physical_block_offset = - (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - K_vec k_vecs[NUM_VECS_PER_THREAD]; - -#pragma unroll - for (int j = 0; j < NUM_VECS_PER_THREAD; j++) { - const cache_t* k_ptr = - k_cache + physical_block_number * kv_block_stride + - kv_head_idx * kv_head_stride + physical_block_offset * x; - const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE; - const int offset1 = (vec_idx * VEC_SIZE) / x; - const int offset2 = (vec_idx * VEC_SIZE) % x; - - if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) { - k_vecs[j] = *reinterpret_cast( - k_ptr + offset1 * BLOCK_SIZE * x + offset2); - } else { - // Vector conversion from Quant_vec to K_vec. - Quant_vec k_vec_quant = *reinterpret_cast( - k_ptr + offset1 * BLOCK_SIZE * x + offset2); - k_vecs[j] = fp8::scaled_convert( - k_vec_quant, *k_scale); - } - } - - // Compute dot product. - // This includes a reduction across the threads in the same thread group. - float qk = scale * Qk_dot::dot( - q_vecs[thread_group_offset], k_vecs); - // Add the ALiBi bias if slopes are given. - qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0; - - if (thread_group_offset == 0) { - // Store the partial reductions to shared memory. - // NOTE(woosuk): It is required to zero out the masked logits. - const bool mask = token_idx >= seq_len; - logits[token_idx - start_token_idx] = mask ? 0.f : qk; - // Update the max value. - qk_max = mask ? qk_max : fmaxf(qk_max, qk); - } - } - } - - // Perform reduction across the threads in the same warp to get the - // max qk value for each "warp" (not across the thread block yet). - // The 0-th thread of each thread group already has its max qk value. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) { - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); - } - if (lane == 0) { - red_smem[warp_idx] = qk_max; - } - __syncthreads(); - - // TODO(woosuk): Refactor this part. - // Get the max qk value for the sequence. - qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX; -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); - } - // Broadcast the max qk value to all threads. - qk_max = VLLM_SHFL_SYNC(qk_max, 0); - - // Get the sum of the exp values. - float exp_sum = 0.f; - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { - float val = __expf(logits[i] - qk_max); - logits[i] = val; - exp_sum += val; - } - exp_sum = block_sum(&red_smem[NUM_WARPS], exp_sum); - - // Compute softmax. - const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f); - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { - logits[i] *= inv_sum; - } - __syncthreads(); - - // If partitioning is enabled, store the max logit and exp_sum. - if (USE_PARTITIONING && thread_idx == 0) { - float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions + partition_idx; - *max_logits_ptr = qk_max; - float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions + partition_idx; - *exp_sums_ptr = exp_sum; - } - - // Each thread will fetch 16 bytes from the value cache at a time. - constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE); - using V_vec = typename Vec::Type; - using L_vec = typename Vec::Type; - using V_quant_vec = typename Vec::Type; - using Float_L_vec = typename FloatVec::Type; - - constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE; - constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW; - constexpr int NUM_ROWS_PER_THREAD = - DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER); - - // NOTE(woosuk): We use FP32 for the accumulator for better accuracy. - float accs[NUM_ROWS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - accs[i] = 0.f; - } - - scalar_t zero_value; - zero(zero_value); - for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; - block_idx += NUM_WARPS) { - // NOTE(woosuk): The block number is stored in int32. However, we cast it to - // int64 because int32 can lead to overflow when this variable is multiplied - // by large numbers (e.g., kv_block_stride). - // For blocksparse attention: skip computation on blocks that are not - // attended - if constexpr (IS_BLOCK_SPARSE) { - int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; - if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) && - !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) { - continue; - } - } - const int64_t physical_block_number = - static_cast(block_table[block_idx]); - const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - L_vec logits_vec; - from_float(logits_vec, *reinterpret_cast(logits + token_idx - - start_token_idx)); - - const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride + - kv_head_idx * kv_head_stride; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE) { - const int offset = row_idx * BLOCK_SIZE + physical_block_offset; - V_vec v_vec; - - if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) { - v_vec = *reinterpret_cast(v_ptr + offset); - } else { - V_quant_vec v_quant_vec = - *reinterpret_cast(v_ptr + offset); - // Vector conversion from V_quant_vec to V_vec. - v_vec = fp8::scaled_convert(v_quant_vec, - *v_scale); - } - if (block_idx == num_seq_blocks - 1) { - // NOTE(woosuk): When v_vec contains the tokens that are out of the - // context, we should explicitly zero out the values since they may - // contain NaNs. See - // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472 - scalar_t* v_vec_ptr = reinterpret_cast(&v_vec); -#pragma unroll - for (int j = 0; j < V_VEC_SIZE; j++) { - v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value; - } - } - accs[i] += dot(logits_vec, v_vec); - } - } - } - - // Perform reduction within each warp. -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - float acc = accs[i]; -#pragma unroll - for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) { - acc += VLLM_SHFL_XOR_SYNC(acc, mask); - } - accs[i] = acc; - } - - // NOTE(woosuk): A barrier is required because the shared memory space for - // logits is reused for the output. - __syncthreads(); - - // Perform reduction across warps. - float* out_smem = reinterpret_cast(shared_mem); -#pragma unroll - for (int i = NUM_WARPS; i > 1; i /= 2) { - int mid = i / 2; - // Upper warps write to shared memory. - if (warp_idx >= mid && warp_idx < i) { - float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - dst[row_idx] = accs[i]; - } - } - } - __syncthreads(); - - // Lower warps update the output. - if (warp_idx < mid) { - const float* src = &out_smem[warp_idx * HEAD_SIZE]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - accs[i] += src[row_idx]; - } - } - } - __syncthreads(); - } - - // Write the final output. - if (warp_idx == 0) { - scalar_t* out_ptr = - out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - from_float(*(out_ptr + row_idx), accs[i]); - } - } - } -} - -// Grid: (num_heads, num_seqs, 1). -template -__global__ void paged_attention_v1_kernel( - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - paged_attention_kernel( - /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache, - v_cache, num_kv_heads, scale, block_tables, seq_lens, - max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride, - kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks, - blocksparse_vert_stride, blocksparse_block_size, - blocksparse_head_sliding_step); -} - -// Grid: (num_heads, num_seqs, max_num_partitions). -template -__global__ void paged_attention_v2_kernel( - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - paged_attention_kernel( - exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale, - block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride, - kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank, - blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size, - blocksparse_head_sliding_step); -} - -// Grid: (num_heads, num_seqs). -template -__global__ void paged_attention_v2_reduce_kernel( - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, - // max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_partitions) { - const int num_heads = gridDim.x; - const int head_idx = blockIdx.x; - const int seq_idx = blockIdx.y; - const int seq_len = seq_lens[seq_idx]; - const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); - if (num_partitions == 1) { - // No need to reduce. Only copy tmp_out to out. - scalar_t* out_ptr = - out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE; - for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) { - out_ptr[i] = tmp_out_ptr[i]; - } - // Terminate the thread block. - return; - } - - constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE; - const int warp_idx = threadIdx.x / WARP_SIZE; - const int lane = threadIdx.x % WARP_SIZE; - - // Size: 2 * num_partitions. - extern __shared__ char shared_mem[]; - // Workspace for reduction. - __shared__ float red_smem[2 * NUM_WARPS]; - - // Load max logits to shared memory. - float* shared_max_logits = reinterpret_cast(shared_mem); - const float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - float max_logit = -FLT_MAX; - for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) { - const float l = max_logits_ptr[i]; - shared_max_logits[i] = l; - max_logit = fmaxf(max_logit, l); - } - __syncthreads(); - - // Get the global max logit. - // Reduce within the warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask)); - } - if (lane == 0) { - red_smem[warp_idx] = max_logit; - } - __syncthreads(); - // Reduce across warps. - max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX; -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask)); - } - // Broadcast the max value to all threads. - max_logit = VLLM_SHFL_SYNC(max_logit, 0); - - // Load rescaled exp sums to shared memory. - float* shared_exp_sums = - reinterpret_cast(shared_mem + sizeof(float) * num_partitions); - const float* exp_sums_ptr = exp_sums + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - float global_exp_sum = 0.0f; - for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) { - float l = shared_max_logits[i]; - float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit); - global_exp_sum += rescaled_exp_sum; - shared_exp_sums[i] = rescaled_exp_sum; - } - __syncthreads(); - global_exp_sum = block_sum(&red_smem[NUM_WARPS], global_exp_sum); - const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f); - - // Aggregate tmp_out to out. - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE; - scalar_t* out_ptr = - out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; -#pragma unroll - for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) { - float acc = 0.0f; - for (int j = 0; j < num_partitions; ++j) { - acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] * - inv_global_exp_sum; - } - from_float(out_ptr[i], acc); - } -} - -} // namespace vllm - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu index 55d75383476..150e3246281 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu @@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options( StrideQ stride_Q_pe = cute::make_tuple( static_cast(q_pe.stride(1)), _1{}, static_cast(q_pe.stride(0))); + // Read the token and page strides from the cache tensor instead of assuming + // packed pages, so strided views (e.g. per-layer views into a cross-layer + // block-major cache) are addressed correctly. StrideK stride_C = cute::make_tuple( - static_cast(0 + D_latent + D_rope), _1{}, static_cast(page_size * (D_latent + D_rope))); + static_cast(kv_c_and_k_pe_cache.stride(1)), _1{}, + static_cast(kv_c_and_k_pe_cache.stride(0))); StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq); StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H); StrideO stride_O = cute::make_tuple(static_cast(0 + D_latent), _1{}, static_cast(0 + H * D_latent)); @@ -268,9 +272,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba using TileShapeD = typename MlaSm100Type::TileShapeD; arguments.problem_shape = cute::make_tuple(TileShapeH{}, static_cast(max_seq_len), TileShapeD{}, static_cast(num_batches)); - // Assumes device 0 when getting sm_count. - arguments.hw_info.sm_count = - sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count; + if (sm_count <= 0) { + int current_device = 0; + cudaGetDevice(¤t_device); + arguments.hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device); + } else { + arguments.hw_info.sm_count = sm_count; + } arguments.split_kv = static_cast(num_kv_splits); MlaSm100Type::Fmha::set_split_kv(arguments); diff --git a/csrc/libtorch_stable/attention/paged_attention_v1.cu b/csrc/libtorch_stable/attention/paged_attention_v1.cu deleted file mode 100644 index 8fa41791593..00000000000 --- a/csrc/libtorch_stable/attention/paged_attention_v1.cu +++ /dev/null @@ -1,190 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, 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 "../torch_utils.h" -#include "attention_kernels.cuh" -#include "../../cuda_compat.h" - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE) \ - VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( \ - ((void*)vllm::paged_attention_v1_kernel), \ - shared_mem_size); \ - vllm::paged_attention_v1_kernel \ - <<>>( \ - out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \ - scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \ - alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \ - k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks, \ - blocksparse_vert_stride, blocksparse_block_size, \ - blocksparse_head_sliding_step); - -// TODO(woosuk): Tune NUM_THREADS. -template -void paged_attention_v1_launcher( - torch::stable::Tensor& out, torch::stable::Tensor& query, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - int num_kv_heads, float scale, torch::stable::Tensor& block_tables, - torch::stable::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, - const int tp_rank, const int blocksparse_local_blocks, - const int blocksparse_vert_stride, const int blocksparse_block_size, - const int blocksparse_head_sliding_step) { - int num_seqs = query.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; - - T* out_ptr = reinterpret_cast(out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.mutable_data_ptr(); - int* seq_lens_ptr = seq_lens.mutable_data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - - const int NUM_WARPS = NUM_THREADS / WARP_SIZE; - int padded_max_seq_len = - DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE; - int logits_size = padded_max_seq_len * sizeof(float); - int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); - // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len - // Keep that in sync with the logic here! - int shared_mem_size = std::max(logits_size, outputs_size); - - dim3 grid(num_heads, num_seqs, 1); - dim3 block(NUM_THREADS); - const torch::stable::accelerator::DeviceGuard device_guard( - query.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - switch (head_size) { - // NOTE(woosuk): To reduce the compilation time, we only compile for the - // head sizes that we use in the model. However, we can easily extend this - // to support any head size which is a multiple of 16. - case 32: - LAUNCH_PAGED_ATTENTION_V1(32); - break; - case 64: - LAUNCH_PAGED_ATTENTION_V1(64); - break; - case 80: - LAUNCH_PAGED_ATTENTION_V1(80); - break; - case 96: - LAUNCH_PAGED_ATTENTION_V1(96); - break; - case 112: - LAUNCH_PAGED_ATTENTION_V1(112); - break; - case 120: - LAUNCH_PAGED_ATTENTION_V1(120); - break; - case 128: - LAUNCH_PAGED_ATTENTION_V1(128); - break; - case 192: - LAUNCH_PAGED_ATTENTION_V1(192); - break; - case 256: - LAUNCH_PAGED_ATTENTION_V1(256); - break; - default: - STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); - break; - } -} - -#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ - paged_attention_v1_launcher( \ - out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \ - seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank, \ - blocksparse_local_blocks, blocksparse_vert_stride, \ - blocksparse_block_size, blocksparse_head_sliding_step); - -#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ - if (is_block_sparse) { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ - } else { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ - } - -// NOTE(woosuk): To reduce the compilation time, we omitted block sizes -// 1, 2, 4, 64, 128, 256. -#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -void paged_attention_v1( - torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, // [num_heads] - double scale, - torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::stable::Tensor& seq_lens, // [num_seqs] - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step) { - const bool is_block_sparse = (blocksparse_vert_stride > 1); - - DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, - CALL_V1_LAUNCHER_BLOCK_SIZE) -} - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/attention/paged_attention_v2.cu b/csrc/libtorch_stable/attention/paged_attention_v2.cu deleted file mode 100644 index 4e8e56ae05c..00000000000 --- a/csrc/libtorch_stable/attention/paged_attention_v2.cu +++ /dev/null @@ -1,202 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, 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 "../torch_utils.h" -#include "attention_kernels.cuh" -#include "../../cuda_compat.h" - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE) \ - vllm::paged_attention_v2_kernel \ - <<>>( \ - exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \ - value_cache_ptr, num_kv_heads, scale, block_tables_ptr, \ - seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, \ - kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank, \ - blocksparse_local_blocks, blocksparse_vert_stride, \ - blocksparse_block_size, blocksparse_head_sliding_step); \ - vllm::paged_attention_v2_reduce_kernel \ - <<>>( \ - out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \ - max_num_partitions); - -template -void paged_attention_v2_launcher( - torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, - torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, - torch::stable::Tensor& query, torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, int num_kv_heads, float scale, - torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, - int max_seq_len, const std::optional& alibi_slopes, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, - const int tp_rank, const int blocksparse_local_blocks, - const int blocksparse_vert_stride, const int blocksparse_block_size, - const int blocksparse_head_sliding_step) { - int num_seqs = query.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; - - T* out_ptr = reinterpret_cast(out.data_ptr()); - float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); - float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); - T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.mutable_data_ptr(); - int* seq_lens_ptr = seq_lens.mutable_data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - - const int NUM_WARPS = NUM_THREADS / WARP_SIZE; - int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); - int logits_size = PARTITION_SIZE * sizeof(float); - int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); - - // For paged attention v2 kernel. - dim3 grid(num_heads, num_seqs, max_num_partitions); - int shared_mem_size = std::max(logits_size, outputs_size); - // For paged attention v2 reduce kernel. - dim3 reduce_grid(num_heads, num_seqs); - int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float); - - dim3 block(NUM_THREADS); - const torch::stable::accelerator::DeviceGuard device_guard( - query.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - switch (head_size) { - // NOTE(woosuk): To reduce the compilation time, we only compile for the - // head sizes that we use in the model. However, we can easily extend this - // to support any head size which is a multiple of 16. - case 32: - LAUNCH_PAGED_ATTENTION_V2(32); - break; - case 64: - LAUNCH_PAGED_ATTENTION_V2(64); - break; - case 80: - LAUNCH_PAGED_ATTENTION_V2(80); - break; - case 96: - LAUNCH_PAGED_ATTENTION_V2(96); - break; - case 112: - LAUNCH_PAGED_ATTENTION_V2(112); - break; - case 120: - LAUNCH_PAGED_ATTENTION_V2(120); - break; - case 128: - LAUNCH_PAGED_ATTENTION_V2(128); - break; - case 192: - LAUNCH_PAGED_ATTENTION_V2(192); - break; - case 256: - LAUNCH_PAGED_ATTENTION_V2(256); - break; - default: - STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); - break; - } -} - -#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ - paged_attention_v2_launcher( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \ - k_scale, v_scale, tp_rank, blocksparse_local_blocks, \ - blocksparse_vert_stride, blocksparse_block_size, \ - blocksparse_head_sliding_step); - -#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ - if (is_block_sparse) { \ - CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ - } else { \ - CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ - } - -// NOTE(woosuk): To reduce the compilation time, we omitted block sizes -// 1, 2, 4, 64, 128, 256. -#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -void paged_attention_v2( - torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - exp_sums, // [num_seqs, num_heads, max_num_partitions] - torch::stable::Tensor& - max_logits, // [num_seqs, num_heads, max_num_partitions] - torch::stable::Tensor& - tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] - torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, // [num_heads] - double scale, - torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::stable::Tensor& seq_lens, // [num_seqs] - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step) { - const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, - CALL_V2_LAUNCHER_BLOCK_SIZE) -} - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index eac93ac9a9f..a1ac81cb10a 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -127,7 +127,12 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, return reinterpret_cast(fn_ptr); }(); - if (batch_fn != nullptr) { + // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 / + // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy + // fallback below, which is correct on any stream. Real and per-thread-default + // streams take the batch fast path. + const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy; + if (batch_fn != nullptr && usable_stream) { CUmemcpyAttributes attr = {}; // ANY lets the DMA engine prefetch source bytes out of stream order, // which is only safe when no GPU stream is concurrently writing the @@ -549,7 +554,7 @@ __global__ void indexer_k_quant_and_cache_kernel( const int head_dim, // dimension of each head const int quant_block_size, // quantization block size const int cache_block_size, // cache block size - const int cache_stride, // stride for each token in kv_cache + const int64_t cache_block_stride, // stride for each block in kv_cache const bool use_ue8m0 // use ue8m0 scale format ) { @@ -590,16 +595,15 @@ __global__ void indexer_k_quant_and_cache_kernel( scale = exp2f(ceilf(log2f(scale))); } - const int64_t dst_offset = block_idx * cache_block_size * cache_stride + - block_offset * head_dim + head_dim_idx; + const int64_t dst_offset = + block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx; for (int i = 0; i < VEC_SIZE; i++) { kv_cache[dst_offset + i] = fp8::scaled_convert(k_val_ptr[i], scale); } if (threadIdx.x == 0) { const int64_t dst_scale_idx = - block_idx * cache_block_size * cache_stride + - cache_block_size * head_dim + + block_idx * cache_block_stride + cache_block_size * head_dim + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; } @@ -1452,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache( reinterpret_cast(k.data_ptr()), \ reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ - cache_block_size, cache_stride, use_ue8m0); + cache_block_size, cache_block_stride, use_ue8m0); void indexer_k_quant_and_cache( torch::stable::Tensor& k, // [num_tokens, head_dim] @@ -1463,7 +1467,7 @@ void indexer_k_quant_and_cache( int num_tokens = k.size(0); int head_dim = k.size(1); int cache_block_size = kv_cache.size(1); - int cache_stride = kv_cache.size(2); + int64_t cache_block_stride = kv_cache.stride(0); bool use_ue8m0 = scale_fmt == "ue8m0"; STD_TORCH_CHECK(k.device() == kv_cache.device(), diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/libtorch_stable/cooperative_topk.cu new file mode 100644 index 00000000000..f388a9e6c8e --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cu @@ -0,0 +1,146 @@ +// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer. +// See cooperative_topk.cuh for kernel implementation. + +#include + +#include "torch_utils.h" + +#ifndef USE_ROCM + #include "cooperative_topk.cuh" +namespace ct = vllm::cooperative; +namespace hist4096 = vllm::topk_histogram_4096; +#endif + +#ifndef USE_ROCM +template +void launch_cooperative_cluster(ct::CooperativeTopKParams& params, + size_t smem, cudaStream_t stream) { + auto kernel = []() { + if constexpr (CS == 16) { + return &ct::cooperative_topk_cs16; + } else if constexpr (CS == 8) { + return &ct::cooperative_topk_cs8; + } else { + static_assert(CS == 4, "unsupported cooperative_topk cluster size"); + return &ct::cooperative_topk_cs4; + } + }(); + if constexpr (CS > 8) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, + 1); + } + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(params.num_rows, CS); + cfg.blockDim = dim3(hist4096::kBlockSize); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim = {1, CS, 1}; + cfg.numAttrs = 1; + cfg.attrs = attrs; + cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params); + STD_TORCH_CHECK(err == cudaSuccess, + "cooperative_topk launch failed: ", cudaGetErrorString(err)); +} + +template +void launch_cooperative_topk_impl(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, + int64_t max_seq_len) { + (void)max_seq_len; // Kept for signature parity with persistent_topk. + const int64_t num_rows = logits.size(0); + const cudaStream_t stream = get_current_cuda_stream(); + + const uint32_t stride = static_cast(logits.stride(0)); + // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves + // headroom) + STD_TORCH_CHECK( + num_rows <= 32, + "cooperative_topk supports <=32 rows; use persistent_topk for " + "larger batches"); + + STD_TORCH_CHECK(stride % 4 == 0, + "cooperative_topk: stride must be multiple of 4 for TMA " + "alignment, got stride (max_model_len)=", + stride); + + STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + STD_TORCH_CHECK( + workspace.scalar_type() == torch::headeronly::ScalarType::Byte, + "workspace must be uint8"); + + ct::CooperativeTopKParams params; + params.input = logits.const_data_ptr(); + params.output = output.mutable_data_ptr(); + params.lengths = lengths.const_data_ptr(); + params.num_rows = static_cast(num_rows); + params.stride = stride; + params.tie_ws = + reinterpret_cast(workspace.mutable_data_ptr()); + + constexpr uint32_t kTieWsPerRow = + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; + STD_TORCH_CHECK( + workspace.size(0) >= + static_cast(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)), + "workspace too small"); + + const bool supports_cluster16 = get_device_prop()->major >= 10; + if (num_rows <= 4 && supports_cluster16) { + launch_cooperative_cluster(params, ct::kSmemSize8, stream); + } else if (num_rows <= 8) { + launch_cooperative_cluster(params, ct::kSmemSize8, stream); + } else { + launch_cooperative_cluster(params, ct::kSmemSize4, stream); + } +} +#endif // USE_ROCM + +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len) { +#ifndef USE_ROCM + STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); + STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); + STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); + STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float, + "Only float32 supported"); + STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int, + "lengths must be int32"); + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int, + "output must be int32"); + STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, + "lengths must be 1D or 2D"); + STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); + STD_TORCH_CHECK(output.dim() == 2, "output must be 2D"); + const int64_t num_rows = logits.size(0); + STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); + STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, + "output size mismatch"); + STD_TORCH_CHECK( + k == 512 || k == 1024 || k == 2048, + "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k); + + if (k == 512) { + launch_cooperative_topk_impl<512>(logits, lengths, output, workspace, + max_seq_len); + } else if (k == 1024) { + launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace, + max_seq_len); + } else { + launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace, + max_seq_len); + } +#else + STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm"); +#endif +} diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh new file mode 100644 index 00000000000..b43b9b8447d --- /dev/null +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -0,0 +1,593 @@ +/* + * Cooperative TopK kernel for DSA Indexer + */ + +#ifndef COOPERATIVE_TOPK_CUH_ +#define COOPERATIVE_TOPK_CUH_ + +#include +#include +#include +#include +#include +#include +#include + +#include "topk_histogram_4096.cuh" + +namespace vllm { +namespace cooperative { + +namespace hist4096 = topk_histogram_4096; + +constexpr uint32_t kHistBits = 10; +constexpr uint32_t kHistBins = 1 << kHistBits; +constexpr uint32_t kMaxTopK = 2048; + +constexpr uint32_t kElemPerStage = 16; +constexpr uint32_t kSizePerStage = + kElemPerStage * hist4096::kBlockSize; // 16384 + +// CS=4 two-pass path uses two TMA stages as a double buffer. +constexpr uint32_t kStreamingStagesCS4 = 2; +// CS=8/16 fused paths keep all loaded TMA stages resident in smem. +constexpr uint32_t kFusedStagesCS8 = 2; +constexpr uint32_t kFusedStagesCS16 = 2; + +// CS=4 single-pass path +constexpr uint32_t kMaxSinglePassStages = 3; +constexpr uint32_t kMaxSinglePassPerBlock = + kMaxSinglePassStages * kSizePerStage; // 49152 + +template +struct CooperativeTopKParams { + const float* __restrict__ input; + int32_t* __restrict__ output; + const int32_t* __restrict__ lengths; + hist4096::Tie* __restrict__ tie_ws; // per-row tie workspace, see + // kTieWsPerRow + uint32_t num_rows, stride; +}; + +// ============================================================================ +// Cooperative helpers +// ============================================================================ + +// only CS adjacent lanes participate (sub-warp reduce), in opposite to +// warp_reduce_sum_full +template +__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) { +#pragma unroll + for (uint32_t m = N >> 1; m > 0; m >>= 1) + v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32); + return v; +} + +// ============================================================================ +// Helpers +// ============================================================================ + +__device__ __forceinline__ uint32_t extract_coarse_bin(float x) { + return hist4096::extract_coarse_bin_N(x); +} + +__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) { + cuda::ptx::mbarrier_init(a, n); +} +__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) { + while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, a, p)); +} +__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a, + uint32_t t) { + cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed, + cuda::ptx::scope_cta, + cuda::ptx::space_shared, a, t); +} +__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n, + uint64_t* m) { + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d, + s, n, m); +} + +// ============================================================================ +// DSMEM histogram reduce +// ============================================================================ + +template +__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) { + static_assert(kHistBins <= hist4096::kBlockSize); + auto cluster = cooperative_groups::this_cluster(); + cluster.sync(); + const auto tx = threadIdx.x; + const auto rank = blockIdx.y; + constexpr auto kLocal = kHistBins / CS; + const auto off = kLocal * rank; + if (tx < kHistBins) { + const auto addr = &histogram[off + tx / CS]; + const auto src = cluster.map_shared_rank(addr, tx % CS); + *src = warp_reduce_sum_subN(*src); + } + cluster.sync(); +} + +// ============================================================================ +// Find threshold from reduced histogram +// ============================================================================ + +// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened +// before calling this, so warp_sum writes are visible across warps. +// The first internal __syncthreads() is still needed for the warp_sum exchange. +template +__device__ __forceinline__ void find_threshold(uint32_t* histogram, + uint32_t* warp_sum, + uint32_t* counter_gt, + uint32_t* counter_eq, + hist4096::MatchBin* match) { + const auto tx = threadIdx.x; + const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize; + const auto value = tx < kHistBins ? histogram[tx] : 0; + const auto winc = hist4096::warp_inclusive_sum(li, value); + if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc; + __syncthreads(); + const auto tmp = warp_sum[li]; + const auto total = hist4096::warp_reduce_sum_full(tmp); + auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc; + const auto above = total - pfx; + if (tx < kHistBins && above < TopK && above + value >= TopK) { + *counter_gt = *counter_eq = 0; + *match = {.bin = tx, .above_count = above, .equal_count = value}; + } + __syncthreads(); +} + +// Streams data through shared memory in chunks, processing each chunk before +// loading the next overwrites each buffer after processing it (the epilogue +// prefetch loads the next chunk into the same slot) +template +__device__ void tma_stream_pass(const float* scores, uint32_t length, + uint32_t thr_bin, int32_t* indices, + uint32_t* phases, SmemType* smem) { + const auto tx = threadIdx.x; + const auto lane = tx % hist4096::kWarpSize; + const auto ni = + (length + kSizePerStage - 1) / kSizePerStage; // total stages needed + const auto la = + (length + 3u) & ~3u; // length rounded up to float4 (TMA alignment) + const auto pass = + kIsScatter ? 1 : 0; // barrier dim: [0] for histogram, [1] for scatter + + // Prologue: issue initial TMA loads - prefill the pipeline + if (tx == 0) { +#pragma unroll + for (uint32_t i = 0; i < kStages; i++) { + if (i >= ni) { + break; + } + const auto o = i * kSizePerStage; + const auto sz = min(kSizePerStage, la - o) * sizeof(float); + tma_load(smem->score_buffer[i], scores + o, sz, + &smem->barrier[pass][i]); // cp.async.bulk is non-blocking + mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz); + } + } + + // Main loop: process stages + for (uint32_t it = 0; it < ni; it++) { + const auto b = it % kStages; // which buffer slot (0 or 1) + const auto o = it * kSizePerStage; + const auto sz = min(kSizePerStage, length - o); + + if (lane == 0) { + mbarrier_wait(&smem->barrier[pass][b], + phases[b] & 1); // wait for the data + } + phases[b]++; // advances the phase for next time this slot is reused + __syncwarp(); + +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i++) { + const auto li = tx + i * hist4096::kBlockSize; + if (li >= sz) { + break; + } + const auto sc = smem->score_buffer[b][li]; + const auto bn = hist4096::extract_coarse_bin_N(sc); + if constexpr (kIsScatter) { // compile-time branch + // Scatter pass: place above-threshold and collect ties + const auto gi = o + li; + if (bn > thr_bin) { + indices[atomicAdd(&smem->counter_gt, 1)] = gi; + } else if (bn == thr_bin) { + const auto p = atomicAdd(&smem->counter_eq, 1); + if (p < hist4096::kMaxTies) { + smem->tie_buffer[p] = {gi, sc}; + } + } + } else { + // Histogram pass: just count + atomicAdd(&smem->histogram[bn], 1); + } + } + __syncthreads(); // ensures all threads finished processing their buffer + // before next TMA load + + // Epilogue: issue next TMA load + if (tx == 0 && it + kStages < ni) { + const auto no = (it + kStages) * kSizePerStage; + const auto nsz = min(kSizePerStage, la - no) * sizeof(float); + tma_load(smem->score_buffer[b], scores + no, nsz, + &smem->barrier[pass][b]); + mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz); + } + } +} + +// ============================================================================ +// Fused path: single TMA pass, rescan smem for scatter +// ============================================================================ + +// Fused shared memory layout for cluster cooperative paths. +// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass +// (CS=4). +template +struct SmemFused { + uint64_t barrier[kPasses][kStages]; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + alignas(128) hist4096::MatchBin match; + uint32_t warp_sum[hist4096::kNumWarps]; + union { + uint32_t histogram[kHistBins]; + hist4096::Tie tie_buffer[kMaxTopK]; + }; + alignas(128) float score_buffer[kStages][kSizePerStage]; +}; + +using Smem8 = SmemFused; +using Smem16 = SmemFused; +using Smem4 = SmemFused; +using SmemSinglePass = SmemFused; + +// Cluster-cooperative large path. +// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan +// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram +// then scatter). +template +__device__ void large_topk(const float* __restrict__ row_input, + int32_t* __restrict__ row_output, uint32_t seq_len, + uint32_t* phases, hist4096::Tie* tie_ws) { + const auto rank = blockIdx.y; // this block's position in cluster + const auto tx = threadIdx.x; + const auto lane = tx % hist4096::kWarpSize; + + extern __shared__ uint8_t smem_raw[]; + auto* smem = reinterpret_cast(smem_raw); + int32_t* s_topk = reinterpret_cast(smem_raw + sizeof(SmemType)); + + // Partition row across cluster ranks + constexpr uint32_t kAlign = 4; + const auto units = + (seq_len + kAlign - 1) / kAlign; // float4-aligned element count + const auto base = units / CS, extra = units % CS; // elements per block + const auto lu = base + (rank < extra ? 1u : 0u); // remainder blocks + const auto ou = + rank * base + min(rank, extra); // this block's count (load-balanced) + const auto my_start = ou * kAlign; // global start offset + const auto my_len = min(my_start + lu * kAlign, seq_len) - + my_start; // actual length of this block + const auto num_iters = + (my_len + kSizePerStage - 1) / kSizePerStage; // TMA stages needed + const auto len_aligned = (my_len + 3u) & ~3u; + + if constexpr (kFused) { + // Fused init + TMA prologue + if (tx < kHistBins) { + smem->histogram[tx] = 0; // all threads zero histogram + } + if (tx == 0) { // thread 0 issues TMA - then all threads continue working + // until mbarrier sync + smem->counter_gt = 0; + smem->counter_eq = 0; + for (uint32_t i = 0; i < num_iters; i++) { + const auto off = i * kSizePerStage; + const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float); + tma_load(smem->score_buffer[i], row_input + my_start + off, sz, + &smem->barrier[0][i]); // cp.async.bulk of size kSizePerStage + // × sizeof(float) + mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz); + } + } + __syncthreads(); + + // Histogram build. ILP unroll-by-2, no inter-stage sync + for (uint32_t iter = 0; iter < num_iters; iter++) { + const auto off = iter * kSizePerStage; + const auto sz = min(kSizePerStage, my_len - off); + if (lane == 0) { + mbarrier_wait(&smem->barrier[0][iter], + phases[iter] & 1); // wait for TMA + } + phases[iter]++; + __syncwarp(); +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i += 2) { + const auto li0 = tx + i * hist4096::kBlockSize; + const auto li1 = tx + (i + 1) * hist4096::kBlockSize; + if (li0 >= sz) { + break; + } + const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]); + if (li1 < sz) { + const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]); + atomicAdd(&smem->histogram[b0], 1); + atomicAdd(&smem->histogram[b1], 1); + } else { + atomicAdd(&smem->histogram[b0], 1); + } + } + } + } else { + // Twopass: init then stream histogram pass + if (tx < kHistBins) { + smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + __syncthreads(); + tma_stream_pass( + row_input + my_start, my_len, 0, nullptr, phases, smem); + } + + // DSMEM all-reduce + find threshold + dsmem_hist_reduce( + smem->histogram); // each block histogram is summed across all CS blocks + find_threshold(smem->histogram, smem->warp_sum, &smem->counter_gt, + &smem->counter_eq, &smem->match); + + const auto thr = smem->match.bin; + + if constexpr (kFused) { + // Fused scatter: rescan score_buffer (still in smem) + for (uint32_t iter = 0; iter < num_iters; iter++) { + const auto off = iter * kSizePerStage; + const auto sz = min(kSizePerStage, my_len - off); +#pragma unroll + for (uint32_t i = 0; i < kElemPerStage; i++) { + const auto li = tx + i * hist4096::kBlockSize; + if (li >= sz) { + break; + } + const auto score = smem->score_buffer[iter][li]; // still in smem + const auto bin = extract_coarse_bin(score); + const auto gidx = off + li; + if (bin > thr) { + s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx; // above -> s_topk + } else if (bin == thr) { + const auto p = atomicAdd(&smem->counter_eq, + 1); // equal -> ties (later refinement) + if (p < hist4096::kMaxTies) { + smem->tie_buffer[p] = {gidx, score}; + } + } + } + } + __syncthreads(); + } else { + // Twopass scatter: re-stream data via TMA + uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0}; + tma_stream_pass( + row_input + my_start, my_len, thr, s_topk, scatter_phases, smem); + } + + // Output collection via DSMEM prefix sum + constexpr uint32_t kAboveBits = 16; + constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1; + static_assert(kAboveMask >= TopK); + static_assert(kAboveMask >= kMaxSinglePassPerBlock, + "kAboveBits must cover max per-block element count"); + + const uint32_t la = smem->counter_gt; + const uint32_t le_full = smem->counter_eq; + const uint32_t le = + min(le_full, hist4096::kMaxTies); // written smem tie_buffer entries + + __shared__ uint32_t s_local_counts[CS]; + __shared__ uint32_t s_prefix_packed; + __shared__ uint32_t s_total_above, s_total_equal; + + auto cluster = cooperative_groups::this_cluster(); + if (tx < CS) { + // Pack written tie counts into 32-bit: (equal << 16) | above. + // `le_full` may exceed the per-block tie buffer cap; using it here creates + // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots. + const uint32_t packed = (le << kAboveBits) | la; + const auto dst = cluster.map_shared_rank(s_local_counts, tx); + dst[rank] = packed; // write my count to every block's s_local_counts[rank] + } + cluster.sync(); + + // Thread 0 computes serial prefix sum + if (tx == 0) { + uint32_t prefix = 0, ta = 0, te = 0; + for (uint32_t i = 0; i < CS; i++) { + if (i == rank) { + s_prefix_packed = prefix; // my prefix + } + ta += s_local_counts[i] & kAboveMask; // total above + te += s_local_counts[i] >> kAboveBits; // total equal + prefix += s_local_counts[i]; + } + s_total_above = ta; + s_total_equal = te; + } + __syncthreads(); + + const uint32_t prefix_above = s_prefix_packed & kAboveMask; + const uint32_t prefix_equal = s_prefix_packed >> kAboveBits; + + // Write to global output + for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) { + // indices are placed contiguously starting at prefix_above + row_output[prefix_above + i] = + s_topk[i] + my_start; // my_start: block-local -> row-global index + } + for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) { + const auto t = smem->tie_buffer[i]; + uint32_t p = s_total_above + prefix_equal + i; + if (p < TopK) { + row_output[p] = t.idx + my_start; + } + uint32_t tp = prefix_equal + i; + if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) { + tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score}; + } + } + + // Tie refinement + cooperative_groups::this_cluster().sync(); + if (rank != 0) { // only rank 0 does tie refinement + return; + } + if (s_total_above + s_total_equal <= TopK) { // no ties to refine + return; + } + + // Tie-breaking uses FP32 (4-round radix sort) + if constexpr (TopK <= hist4096::kBlockSize) { + // copy ties from tie_ws back to smem, then refine + const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies); + // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie) + for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) { + smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score}; + } + __syncthreads(); + hist4096::tie_handle(smem->tie_buffer, num_ties, s_total_above, + row_output, smem); + } else { + // TopK=2048: process directly from tie_ws (GMEM) + const uint32_t num_ties = min(s_total_equal, static_cast(TopK)); + hist4096::tie_handle_large(tie_ws, num_ties, s_total_above, + row_output, smem); + } +} + +// ============================================================================ +// Adapted from https://github.com/sgl-project/sglang/pull/23600 +// sgl-project/sglang +// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/) +// ============================================================================ + +template +__device__ void cooperative_topk_body(CooperativeTopKParams params) { + const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x; + const auto sl = params.lengths[row]; + int32_t* out = params.output + row * TopK; + const float* in = params.input + row * params.stride; + + // Trivial: seq_len <= TopK + if (sl <= static_cast(TopK)) { + if (rank == 0) { + for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) { + out[i] = (i < static_cast(sl)) ? static_cast(i) : -1; + } + } + return; + } + + // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF + if (sl <= static_cast(hist4096::kHist4096MaxLen)) { + if (rank == 0) { + extern __shared__ uint8_t sr[]; + hist4096::histogram_4096_topk( + in, out, sl, sr); // 4096-bin (12-bit) histogram + } + return; + } + + // Large path: init mbarriers + state, then dispatch fused or twopass + const uint32_t per_block = + (params.stride + CS - 1) / CS; // how many elements per block + constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 + : kMaxSinglePassStages) * + kSizePerStage; + const bool use_singlepass = + per_block <= + kFusedMax; // single pass or TMA streaming: histogram+scatter + + // Select smem type and stage count at compile time based on CS + constexpr uint32_t kFusedStages = (CS == 16) ? kFusedStagesCS16 + : (CS == 8) ? kFusedStagesCS8 + : kMaxSinglePassStages; + using FusedSmem = SmemFused; + + extern __shared__ uint8_t sr[]; + + constexpr uint32_t kTieWsPerRow = + TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK; + hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow; + + if (use_singlepass) { + auto* smem = reinterpret_cast(sr); + const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage; + if (tx < sp_stages) { + mbarrier_init(&smem->barrier[0][tx], + 1); // init 1 barrier per TMA stage - + // signal when async copies complete + } + __syncthreads(); + uint32_t phases[kFusedStages] = + {}; // tracks the parity for mbarrier wait/arrive protocol + large_topk(in, out, sl, phases, row_tie_ws); + } else { + // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass) + auto* smem = reinterpret_cast(sr); + if (tx < 2 * kStreamingStagesCS4) { + mbarrier_init(&smem->barrier[0][tx], + 1); // init 2×2=4 barriers (2 passes × 2 stages) + } + __syncthreads(); + uint32_t hp[kStreamingStagesCS4] = {0, + 0}; // histogram+scatter pass counters + large_topk(in, out, sl, hp, row_tie_ws); + } +} + +template +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 4, 1) + cooperative_topk_cs4(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + +template +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 8, 1) + cooperative_topk_cs8(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + +template +__global__ void __launch_bounds__(hist4096::kBlockSize, 1) + __cluster_dims__(1, 16, 1) + cooperative_topk_cs16(CooperativeTopKParams params) { + cooperative_topk_body(params); +} + +constexpr size_t kSmemSize4_base = sizeof(Smem4); +constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass); +constexpr size_t kSmemSize4 = + (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) + + sizeof(int32_t) * 2048 + 128; +constexpr size_t kSmemSize8 = + sizeof(SmemFused) + sizeof(int32_t) * 2048 + 128; + +} // namespace cooperative + +} // namespace vllm + +#endif // COOPERATIVE_TOPK_CUH_ diff --git a/csrc/core/math.hpp b/csrc/libtorch_stable/core/math.hpp similarity index 100% rename from csrc/core/math.hpp rename to csrc/libtorch_stable/core/math.hpp diff --git a/csrc/cub_helpers.h b/csrc/libtorch_stable/cub_helpers.h similarity index 100% rename from csrc/cub_helpers.h rename to csrc/libtorch_stable/cub_helpers.h diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/libtorch_stable/cuda_vec_utils.cuh index efbb09994d2..ec6e60724e6 100644 --- a/csrc/libtorch_stable/cuda_vec_utils.cuh +++ b/csrc/libtorch_stable/cuda_vec_utils.cuh @@ -21,7 +21,7 @@ // together enable 256-bit (v8.u32) PTX load/store instructions. // Use for PTX instruction selection with architecture fallback paths. #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ - defined(CUDA_VERSION) && CUDA_VERSION >= 12090 + defined(CUDART_VERSION) && CUDART_VERSION >= 12090 #define VLLM_256B_PTX_ENABLED 1 #else #define VLLM_256B_PTX_ENABLED 0 diff --git a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp similarity index 100% rename from csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp rename to csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp diff --git a/csrc/cutlass_extensions/torch_utils.hpp b/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp similarity index 100% rename from csrc/cutlass_extensions/torch_utils.hpp rename to csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp diff --git a/csrc/cutlass_extensions/vllm_collective_builder.cuh b/csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh similarity index 100% rename from csrc/cutlass_extensions/vllm_collective_builder.cuh rename to csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh diff --git a/csrc/cutlass_extensions/vllm_numeric_conversion.cuh b/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh similarity index 100% rename from csrc/cutlass_extensions/vllm_numeric_conversion.cuh rename to csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh diff --git a/csrc/libtorch_stable/dispatch_utils.h b/csrc/libtorch_stable/dispatch_utils.h index e9478236a0e..cd67ac751c4 100644 --- a/csrc/libtorch_stable/dispatch_utils.h +++ b/csrc/libtorch_stable/dispatch_utils.h @@ -30,6 +30,28 @@ THO_DISPATCH_SWITCH(TYPE, NAME, \ VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) + // FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format #ifdef USE_ROCM #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu index bdf749ddfcf..585004c047b 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu @@ -328,7 +328,7 @@ struct GmemLoaderB { __device__ void issue_mainloop() { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #pragma unroll 1 for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { if (need_wait) { @@ -643,7 +643,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( mma_computer.issue_mainloop(); mma_computer.epi(); } - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -733,6 +733,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, output.scalar_type() == torch::headeronly::ScalarType::BFloat16, "Only BFloat16 output dtype is supported"); + const torch::stable::accelerator::DeviceGuard device_guard( + mat_a.get_device_index()); STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu index 04397e0893c..3207082d579 100644 --- a/csrc/libtorch_stable/fp32_router_gemm.cu +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -1,13 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32. +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the +// supported (E, H) pairs listed at the bottom of this file. // Supports bf16 or fp32 activation; weight is always fp32. // Adapted from dsv3_router_gemm_float_out.cu. +// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see +// invokeFp32RouterGemm. #include #include +#include + // --------------------------------------------------------------------------- // Load helpers // --------------------------------------------------------------------------- @@ -73,94 +78,113 @@ __device__ __forceinline__ void load_activation<__nv_bfloat16, 8>( // InputT : type of activation (float or __nv_bfloat16) // Weight is always fp32; output is always fp32. // VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16 -template -__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel( - float* out, InputT const* mat_a, float const* mat_b) { +// Each block computes kEPB expert columns; wider blocks / kEPB > 1 are +// selected per (shape, M) in invokeFp32RouterGemm (B300-tuned, see below). +// kTGroups > 1 splits the tokens across groups of kBlockSize threads within +// the block: all groups scan the same weight K-slices (group 0 misses to +// DRAM, later groups hit L1) so weight traffic stays 1x, while per-thread +// accumulator registers drop by kTGroups (at M=16 the 32 fp32 accumulators +// push the kernel to 128 regs/thread and 1 block/SM). +template +__global__ __launch_bounds__( + kBlockSize* kTGroups, 1) void fp32_router_gemm_kernel(float* out, + InputT const* mat_a, + float const* mat_b) { constexpr int VPT = 16 / sizeof(InputT); constexpr int k_elems_per_k_iteration = VPT * kBlockSize; constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; + static_assert(kHiddenDim % k_elems_per_k_iteration == 0); + static_assert(kNumTokens % kTGroups == 0); constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; + constexpr int kNumWarps = kBlockSize / kWarpSize; // per token group + constexpr int kMG = kNumTokens / kTGroups; // tokens per group - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; + int const e_base = blockIdx.x * kEPB; + int const tid = threadIdx.x % kBlockSize; + int const m0 = (threadIdx.x / kBlockSize) * kMG; int const warpId = tid / kWarpSize; int const laneId = tid % kWarpSize; - float acc[kNumTokens] = {}; - __shared__ float sm_reduction[kNumTokens][kNumWarps]; - - float const* b_col = mat_b + n_idx * kHiddenDim; - - int k_bases[k_iterations]; -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) { - k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; - } + float acc[kMG][kEPB] = {}; + __shared__ float sm_reduction[kNumTokens][kEPB][kNumWarps]; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); + // Fire the PDL trigger right after our own wait instead of at kernel end: + // a gridsync-ing consumer is unaffected (its wait always targets full grid + // completion), while a consumer that reads none of our outputs (e.g. the + // NVFP4 activation quant, which reads the same hidden_states) can launch + // now and fully overlap this kernel's body. + cudaTriggerProgrammaticLaunchCompletion(); #endif +#pragma unroll for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; + int const k_base = ki * k_elems_per_k_iteration + tid * VPT; - float b_float[VPT]; - load_weight(b_col + k_base, b_float); + float b_float[kEPB][VPT]; +#pragma unroll + for (int e = 0; e < kEPB; e++) { + load_weight(mat_b + (e_base + e) * kHiddenDim + k_base, b_float[e]); + } #pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { + for (int m_idx = 0; m_idx < kMG; m_idx++) { float a_float[VPT]; - load_activation(mat_a + m_idx * kHiddenDim + k_base, - a_float); + load_activation( + mat_a + (size_t)(m0 + m_idx) * kHiddenDim + k_base, a_float); #pragma unroll - for (int k = 0; k < VPT; k++) { - acc[m_idx] += a_float[k] * b_float[k]; + for (int e = 0; e < kEPB; e++) { +#pragma unroll + for (int k = 0; k < VPT; k++) { + acc[m_idx][e] += a_float[k] * b_float[e][k]; + } } } } // Warp-level butterfly reduction #pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = acc[m]; - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - if (laneId == 0) sm_reduction[m][warpId] = sum; + for (int m = 0; m < kMG; m++) { +#pragma unroll + for (int e = 0; e < kEPB; e++) { + float sum = acc[m][e]; + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (laneId == 0) sm_reduction[m0 + m][e][warpId] = sum; + } } __syncthreads(); - if (tid == 0) { + // Parallel finalize: one thread per (m, e) output. + for (int idx = threadIdx.x; idx < kNumTokens * kEPB; + idx += kBlockSize * kTGroups) { + int const m = idx / kEPB; + int const e = idx % kEPB; + float final_sum = 0.0f; #pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; -#pragma unroll - for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w]; - out[m * kNumExperts + n_idx] = final_sum; - } + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][e][w]; + out[m * kNumExperts + e_base + e] = final_sum; } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.launch_dependents;"); -#endif } // --------------------------------------------------------------------------- // Launcher // --------------------------------------------------------------------------- -template -void invokeFp32RouterGemm(float* output, InputT const* mat_a, - float const* mat_b, cudaStream_t stream) { - constexpr int kBlockSize = 128; +template +static void launchFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + static_assert(kNumExperts % kEPB == 0); cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; + config.gridDim = kNumExperts / kEPB; + config.blockDim = kBlockSize * kTGroups; config.dynamicSmemBytes = 0; config.stream = stream; cudaLaunchAttribute attrs[1]; @@ -168,56 +192,158 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, attrs[0].val.programmaticStreamSerializationAllowed = 1; config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, - fp32_router_gemm_kernel, - output, mat_a, mat_b); + cudaLaunchKernelEx( + &config, + fp32_router_gemm_kernel, + output, mat_a, mat_b); +} + +static bool isBlackwellFamily() { + static int sm = []() { + int dev = 0, major = 0, minor = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev); + return major * 10 + minor; + }(); + return sm >= 100; +} + +template +void invokeFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + // Geometry tuned on B300 per supported shape, bf16 activation, under a + // production-fidelity harness (CUDA-graph replay, per-layer cold weights). + // GLM-5.2 (E=256, H=6144): + // M <= 4 : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1) + // M in [5, 15] + // or odd : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8)) + // M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04, + // M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at + // 0.97x, so the boundary is 16). + // Only enabled on the Blackwell family where it was validated; Hopper and + // other shapes / fp32 activation keep the legacy geometry. + if constexpr (std::is_same_v && kNumExperts == 256 && + kHiddenDim == 6144) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens <= 4) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 128 && kHiddenDim == 6144) { + // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same + // accumulator register cliffs; B300 sweep: + // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x) + // even M >= 12 : BS=192, EPB=1, 2 token groups (1.59-1.66x at + // M >= 18; re-measured on B300+B200: 192 also wins + // M=12/14 by 5-11%% on both, ties 384 at 16) + // M <= 5 / odd : BS=384, EPB=1 (1.03-1.19x) + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 256 && kHiddenDim == 3072) { + // MiniMax-M2/M2.5. The 3.1MB weight is latency-floor bound at small M + // (legacy already optimal); token groups win only at even M >= 8 + // (1.05-1.17x). EPB crossover measured between 12 and 16. + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 14 && kNumTokens % 2 == 0) { + // M=14 originally measured 0.91x and stayed on legacy; two fresh + // sweeps (B300 dev1 + B200) both put 192/2/tg2 ahead by 3.5-4%%. + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 && + kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } } // --------------------------------------------------------------------------- -// Explicit instantiations: M=1..32, E=256, H=3072, for both input types +// Explicit instantiations: M=1..32, for both input types, for the supported +// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3] +// and (256, 6144) [GLM-5.2]. // --------------------------------------------------------------------------- -#define INSTANTIATE(T, M) \ - template void invokeFp32RouterGemm( \ - float*, T const*, float const*, cudaStream_t); +#define INSTANTIATE(T, M, E, H) \ + template void invokeFp32RouterGemm(float*, T const*, \ + float const*, cudaStream_t); -#define INSTANTIATE_ALL(T) \ - INSTANTIATE(T, 1) \ - INSTANTIATE(T, 2) \ - INSTANTIATE(T, 3) \ - INSTANTIATE(T, 4) \ - INSTANTIATE(T, 5) \ - INSTANTIATE(T, 6) \ - INSTANTIATE(T, 7) \ - INSTANTIATE(T, 8) \ - INSTANTIATE(T, 9) \ - INSTANTIATE(T, 10) \ - INSTANTIATE(T, 11) \ - INSTANTIATE(T, 12) \ - INSTANTIATE(T, 13) \ - INSTANTIATE(T, 14) \ - INSTANTIATE(T, 15) \ - INSTANTIATE(T, 16) \ - INSTANTIATE(T, 17) \ - INSTANTIATE(T, 18) \ - INSTANTIATE(T, 19) \ - INSTANTIATE(T, 20) \ - INSTANTIATE(T, 21) \ - INSTANTIATE(T, 22) \ - INSTANTIATE(T, 23) \ - INSTANTIATE(T, 24) \ - INSTANTIATE(T, 25) \ - INSTANTIATE(T, 26) \ - INSTANTIATE(T, 27) \ - INSTANTIATE(T, 28) \ - INSTANTIATE(T, 29) \ - INSTANTIATE(T, 30) \ - INSTANTIATE(T, 31) \ - INSTANTIATE(T, 32) +#define INSTANTIATE_ALL(T, E, H) \ + INSTANTIATE(T, 1, E, H) \ + INSTANTIATE(T, 2, E, H) \ + INSTANTIATE(T, 3, E, H) \ + INSTANTIATE(T, 4, E, H) \ + INSTANTIATE(T, 5, E, H) \ + INSTANTIATE(T, 6, E, H) \ + INSTANTIATE(T, 7, E, H) \ + INSTANTIATE(T, 8, E, H) \ + INSTANTIATE(T, 9, E, H) \ + INSTANTIATE(T, 10, E, H) \ + INSTANTIATE(T, 11, E, H) \ + INSTANTIATE(T, 12, E, H) \ + INSTANTIATE(T, 13, E, H) \ + INSTANTIATE(T, 14, E, H) \ + INSTANTIATE(T, 15, E, H) \ + INSTANTIATE(T, 16, E, H) \ + INSTANTIATE(T, 17, E, H) \ + INSTANTIATE(T, 18, E, H) \ + INSTANTIATE(T, 19, E, H) \ + INSTANTIATE(T, 20, E, H) \ + INSTANTIATE(T, 21, E, H) \ + INSTANTIATE(T, 22, E, H) \ + INSTANTIATE(T, 23, E, H) \ + INSTANTIATE(T, 24, E, H) \ + INSTANTIATE(T, 25, E, H) \ + INSTANTIATE(T, 26, E, H) \ + INSTANTIATE(T, 27, E, H) \ + INSTANTIATE(T, 28, E, H) \ + INSTANTIATE(T, 29, E, H) \ + INSTANTIATE(T, 30, E, H) \ + INSTANTIATE(T, 31, E, H) \ + INSTANTIATE(T, 32, E, H) -INSTANTIATE_ALL(float) -INSTANTIATE_ALL(__nv_bfloat16) +INSTANTIATE_ALL(float, 256, 3072) +INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) +INSTANTIATE_ALL(float, 128, 6144) +INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) +INSTANTIATE_ALL(float, 256, 6144) +INSTANTIATE_ALL(__nv_bfloat16, 256, 6144) #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu index 4baa740de93..8643e33cbe7 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -22,36 +22,44 @@ inline int getSMVersion() { } // namespace -static constexpr int FP32_NUM_EXPERTS = 256; -static constexpr int FP32_HIDDEN_DIM = 3072; static constexpr int FP32_MAX_TOKENS = 32; +// Supported (hidden_dim, num_experts) pairs (must match the instantiations in +// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3, +// (6144, 256) for GLM-5.2. +static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { + return (hidden_dim == 3072 && num_experts == 256) || + (hidden_dim == 6144 && num_experts == 128) || + (hidden_dim == 6144 && num_experts == 256); +} + // Forward declarations — 4 template params must match fp32_router_gemm.cu template void invokeFp32RouterGemm(float* output, InputT const* mat_a, float const* mat_b, cudaStream_t stream); -// LoopUnroller templated on InputT -template +// LoopUnroller templated on InputT, kNumExperts and kHiddenDim +template struct Fp32LoopUnroller { static void unroll(int num_tokens, float* output, InputT const* mat_a, float const* mat_b, cudaStream_t stream) { if (num_tokens == kBegin) { - invokeFp32RouterGemm( + invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { - Fp32LoopUnroller::unroll(num_tokens, output, - mat_a, mat_b, stream); + Fp32LoopUnroller::unroll(num_tokens, output, mat_a, mat_b, stream); } } }; -template -struct Fp32LoopUnroller { +template +struct Fp32LoopUnroller { static void unroll(int num_tokens, float* output, InputT const* mat_a, float const* mat_b, cudaStream_t stream) { if (num_tokens == kEnd) { - invokeFp32RouterGemm( + invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { throw std::invalid_argument( @@ -60,6 +68,26 @@ struct Fp32LoopUnroller { } }; +// Dispatch over the supported (num_experts, hidden_dim) pairs. +template +void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, + float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + if (num_experts == 256 && hidden_dim == 3072) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 128 && hidden_dim == 6144) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 6144) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else { + throw std::invalid_argument( + "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); + } +} + void fp32_router_gemm( torch::stable::Tensor& output, // [num_tokens, num_experts] torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim] @@ -85,10 +113,10 @@ void fp32_router_gemm( STD_TORCH_CHECK( mat_a.size(1) == mat_b.size(1), "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim"); - STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM, - "fp32_router_gemm: expected hidden_dim=3072"); - STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS, - "fp32_router_gemm: expected num_experts=256"); + STD_TORCH_CHECK( + fp32_router_gemm_supported(hidden_dim, num_experts), + "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " + "(3072, 256), (6144, 128) and (6144, 256)"); STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, "fp32_router_gemm: num_tokens must be in [0, 32]"); STD_TORCH_CHECK( @@ -104,6 +132,8 @@ void fp32_router_gemm( return; } + const torch::stable::accelerator::DeviceGuard device_guard( + mat_a.get_device_index()); STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); @@ -113,12 +143,13 @@ void fp32_router_gemm( if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { auto const* mat_a_ptr = reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); - Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll( - num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream); + dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens, + out_ptr, mat_a_ptr, mat_b_ptr, + stream); } else { auto const* mat_a_ptr = reinterpret_cast(mat_a.data_ptr()); - Fp32LoopUnroller::unroll( - num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream); + dispatchFp32RouterGemm(num_experts, hidden_dim, num_tokens, out_ptr, + mat_a_ptr, mat_b_ptr, stream); } } diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 4d34b4b6b50..7bc435b8e0d 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -18,7 +18,7 @@ * 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 + * FP8_MAX = 224.0f on ROCm FNUZ / 448.0f on OCP * 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) @@ -61,10 +61,11 @@ #ifdef USE_ROCM // ROCm-compatible FP8 conversion helpers __device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { - #if defined(HIP_FP8_TYPE_OCP) - __hip_fp8_e4m3 fp8_val(val); - #else + // gfx942 uses FNUZ FP8; other ROCm targets use OCP E4M3. + #if defined(__gfx942__) __hip_fp8_e4m3_fnuz fp8_val(val); + #else + __hip_fp8_e4m3 fp8_val(val); #endif return reinterpret_cast(fp8_val); } @@ -90,7 +91,13 @@ 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 +// FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw +// 240.0) to match the rest of vLLM's FNUZ pipeline. +#if defined(USE_ROCM) && defined(__gfx942__) +constexpr float kFp8Max = 224.0f; +#else constexpr float kFp8Max = 448.0f; +#endif #ifndef USE_ROCM // When num_tokens is less than this threshold, diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu new file mode 100644 index 00000000000..d8460f032eb --- /dev/null +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -0,0 +1,826 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Horizontally-fused MiniMax-M3 attention pre-processing kernel. + * + * Replaces the per-token Python sequence in + * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``: + * + * q = q_norm(q); k = k_norm(k); q, k = rotary_emb(pos, q, k) + * index_q = index_q_norm(index_q); index_k = index_k_norm(index_k) + * index_q, index_k = rotary_emb(pos, index_q, index_k) + * _insert_kv(k, v, index_k) + * + * All branches share head_dim=128 and the *same* partial-NeoX RoPE table + * (``rotary_dim`` rotated, the trailing dims pass through). The four norms + * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with + * independent weights. + * + * Everything lives in a single fused ``qkv`` tensor. The sparse layer's + * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token:: + * + * [ q | k | v | index_q | index_k ] (the "5 results") + * + * while the dense layer emits just ``[ q | k | v ]``. The kernel reads the + * index branch straight out of that packed row -- no separate index tensors. + * + * One kernel, one grid; each warp owns one (token, head-slot) pair. Slot + * enumeration per token: + * [0, nq) Q heads -> norm(q_w) + RoPE, write + * qkv [nq, nq+nkv) K heads -> norm(k_w) + RoPE, write + * qkv + * (+ insert into key cache) + * [nq+nkv, nq+2*nkv) V heads -> insert into value cache + * IQ heads (niq) -> norm(iq_w) + RoPE, write iq + * IK (1) -> norm(ik_w) + RoPE + * (+ insert into index cache) + * + * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the + * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128. + * + * Dense vs sparse row layout and index-branch processing are separate template + * choices. Skip-index-topk reuse layers still have sparse rows and insert main + * K/V cache entries, but compile away index_q/index_k work and index-cache + * writes. + * + * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused + * ``qkv`` tensor. Caches (bf16) are scatter-written by slot. + */ + +#include +#include +#include + +#include "torch_utils.h" + +#include "../cuda_compat.h" +#include "type_convert.cuh" +#include "../attention/dtype_fp8.cuh" +#include "dispatch_utils.h" + +#ifdef USE_ROCM + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#else + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#endif + +// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs. +#ifndef USE_ROCM + #include +#else + #include +#endif + +#ifndef FINAL_MASK + #ifdef USE_ROCM + #define FINAL_MASK 0xffffffffffffffffULL + #else + #define FINAL_MASK 0xffffffffu + #endif +#endif + +#ifdef USE_ROCM +// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4 +// fused kernel). +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { + #if defined(HIP_FP8_TYPE_OCP) + __hip_fp8_e4m3 fp8_val(val); + #else + __hip_fp8_e4m3_fnuz fp8_val(val); + #endif + return reinterpret_cast(fp8_val); +} +#endif + +namespace vllm { +namespace minimax_m3_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} +} // namespace + +// ──────────────────────────────────────────────────────────────────────────── +// Constants (hard-coded for MiniMax-M3-preview). +// ──────────────────────────────────────────────────────────────────────────── +constexpr int kHeadDim = 128; +constexpr int kNumLanes = 32; +constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 4 + +// ──────────────────────────────────────────────────────────────────────────── +// Helpers +// ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warpReduceSum(float val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val += __shfl_xor_sync(FINAL_MASK, val, mask, 32); + } + return val; +} + +// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``), rounded +// back to scalar_t like the materialized unfused norm output, followed by +// partial NeoX RoPE on the leading ``rotary_dim`` dims. Each lane owns +// ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4). +template +__device__ __forceinline__ void normAndRope( + float (&elems)[kElemsPerLane], int const laneId, float const eps, + scalar_t const* __restrict__ weight, // [kHeadDim] or nullptr (no norm) + bool const do_rope, int const rotary_dim, + scalar_t const* __restrict__ cos_ptr, // cos_sin_cache + pos*rotary_dim + bool const apply_norm) { + // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ────────────────── + if (apply_norm) { + float sumsq = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i]; + sumsq = warpReduceSum(sumsq); + float const rms_rcp = rsqrtf(sumsq / static_cast(kHeadDim) + eps); +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + int const dim = laneId * kElemsPerLane + i; + float const w = 1.0f + static_cast(weight[dim]); + elems[i] = elems[i] * rms_rcp * w; + } + } + + // ── Partial NeoX RoPE on dims [0, rotary_dim) ────────────────────────── + // half = rotary_dim/2. Pair (i, i+half) for i in [0, half). Lane L owns + // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the + // first half (own=x[i]) or second half (own=x[i+half]); its partner lives + // ``half/4`` lanes away (XOR with that distance). + if (do_rope) { + int const half = rotary_dim / 2; + int const dim0 = laneId * kElemsPerLane; + bool const in_rope = dim0 < rotary_dim; + int const lane_xor = half / kElemsPerLane; // partner-lane distance + + float partner[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32); + } + if (in_rope) { + bool const first_half = dim0 < half; + int const i_base = first_half ? dim0 : (dim0 - half); // cos/sin index + scalar_t const* sin_ptr = cos_ptr + half; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + float const c = static_cast(cos_ptr[i_base + i]); + float const s = static_cast(sin_ptr[i_base + i]); + if (first_half) { + elems[i] = elems[i] * c - partner[i] * s; + } else { + elems[i] = elems[i] * c + partner[i] * s; + } + } + } + } +} + +// Load 4 contiguous bf16 -> 4 fp32 registers. +template +__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src, + float (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + uint2 v = *reinterpret_cast(src); + auto const* p = + reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + float2 f2 = Converter::convert(p[i]); + elems[2 * i] = f2.x; + elems[2 * i + 1] = f2.y; + } +} + +// Store 4 fp32 registers -> 4 contiguous bf16. +template +__device__ __forceinline__ void storeElems( + scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + uint2 v; + auto* p = reinterpret_cast(&v); +#pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1])); + } + *reinterpret_cast(dst) = v; +} + +// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache +// dtypes use the scaled-convert path with identity scale. +template +__device__ __forceinline__ void storeCacheElems( + cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { + // kAuto means unquantized KV cache here: cache_t == scalar_t, so store the + // model dtype directly. FP8 cache dtypes use the conversion path below. + storeElems(reinterpret_cast(dst), elems); + } else { +#ifdef USE_ROCM + // Match ROCm's model-dtype materialization before FP8 cache conversion. + using Converter = vllm::_typeConvert; + using rounded_t = typename Converter::hip_type; + rounded_t rounded[kElemsPerLane]; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + rounded[i] = Converter::convert(elems[i]); + } + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + dst[i] = fp8::scaled_convert(rounded[i], 1.0f); + } +#else + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + dst[i] = fp8::scaled_convert(elems[i], 1.0f); + } +#endif + } +} + +// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast, +// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale +// (RMSNorm outputs are O(1) and the score path only needs relative block +// ordering). +__device__ __forceinline__ void storeElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + constexpr float kFp8Max = 448.0f; +#ifndef USE_ROCM + __nv_fp8x2_storage_t out2[kElemsPerLane / 2]; + #pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]); + vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max); + vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max); + out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3); + } + *reinterpret_cast(dst) = *reinterpret_cast(out2); +#else + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max); + dst[i] = rocm_cvt_float_to_fp8_e4m3(vv); + } +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Kernel +// ──────────────────────────────────────────────────────────────────────────── +// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block). +// Each warp = one (token, slot). +// +// `kHasIndex`, `kProcessIndex`, and `kInsertKV` are compile-time template +// bools, so branch decisions that distinguish the dense layer from the sparse +// layer (index slots, KV/index inserts, V slots) fold away per instantiation. +// Slots per token: +// Q : nq (always — norm+RoPE) +// K : nkv (always — norm+RoPE; +K-cache insert) +// V : nkv only if kInsertKV (V-cache insert; no warps in dense) +// IQ: niq only if kProcessIndex (norm+RoPE) +// IK: 1 only if kProcessIndex (norm+RoPE; +index-cache insert) +// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx: +// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte). +// kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k]. +// kProcessIndex controls whether this launch actually norms/ropes the index +// branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep +// kHasIndex=true but set kProcessIndex=false. +template +__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( + scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse) + scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr + out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte + scalar_t const* __restrict__ q_norm_w, + scalar_t const* __restrict__ k_norm_w, + scalar_t const* __restrict__ iq_norm_w, + scalar_t const* __restrict__ ik_norm_w, + scalar_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim] + int64_t const* __restrict__ positions, // [N] i64 + int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr + int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr + cache_t* __restrict__ kv_cache, // [nb,nkv,bs,2*128] or nullptr + out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte + float const eps, int const rotary_dim, int const num_tokens, int const nq, + int const nkv, int const niq, int const block_size, + // kv_cache strides (in elements) for logical shape [nb, nkv, bs, 2*128]. + // The content (last) dim is always innermost-contiguous (stride 1), so the + // NHD/HND layout choice is captured by the head/token strides. + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + // _typeConvert is unavailable on pre-Ampere; the M3 kernel only + // runs with bf16/fp16 inputs in practice. Discard the bf16 body there. + if constexpr (std::is_same_v) { + return; + } else { +#endif + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32); + + static_assert(!kProcessIndex || kHasIndex, + "index processing requires sparse row layout"); + + // Slot layout (compile-time gated: dense has neither V nor index slots). + int const v_slots = kInsertKV ? nkv : 0; + int const idx_slots = kProcessIndex ? niq + 1 : 0; + int const slots_per_token = nq + nkv + v_slots + idx_slots; + + int const tokenIdx = globalWarpIdx / slots_per_token; + int const slot = globalWarpIdx % slots_per_token; + if (tokenIdx >= num_tokens) return; + + // Slot boundaries. + int const k_begin = nq; + int const v_begin = nq + nkv; // valid only when kInsertKV + int const iq_begin = nq + nkv + v_slots; // index block start + int const ik_slot = iq_begin + niq; // valid only when kProcessIndex + + bool const isQ = slot < k_begin; + bool const isK = slot >= k_begin && slot < v_begin; + bool isV = false; + if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv; + bool isIQ = false, isIK = false; + if constexpr (kProcessIndex) { + isIQ = slot >= iq_begin && slot < ik_slot; + isIK = slot == ik_slot; + } + + int const dim_base = laneId * kElemsPerLane; + // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse + // layer additionally packs [index_q (niq heads) | index_k (1 head)]. + int const qkv_row = (nq + 2 * nkv + (kHasIndex ? (niq + 1) : 0)) * kHeadDim; + + // ── Resolve source pointer + per-branch parameters. ──────────────────── + scalar_t* row_ptr = nullptr; // in-place output location + scalar_t const* norm_w = nullptr; // nullptr -> skip norm (V) + bool do_rope = true; + int head = 0; // kv head index for inserts + + if (isQ) { + row_ptr = + qkv + static_cast(tokenIdx) * qkv_row + slot * kHeadDim; + norm_w = q_norm_w; + } else if (isK) { + head = slot - k_begin; + row_ptr = + qkv + static_cast(tokenIdx) * qkv_row + slot * kHeadDim; + norm_w = k_norm_w; + } else if (isV) { + // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the + // correct in-tensor offset. + head = slot - v_begin; + row_ptr = + qkv + static_cast(tokenIdx) * qkv_row + slot * kHeadDim; + norm_w = nullptr; // V: no norm, no rope + do_rope = false; + } else if (isIQ) { + // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv. + int const ih = slot - iq_begin; + row_ptr = qkv + static_cast(tokenIdx) * qkv_row + + (nq + 2 * nkv + ih) * kHeadDim; + norm_w = iq_norm_w; + } else if (isIK) { + // Single shared index key at (nq+2*nkv+niq)*128. + row_ptr = qkv + static_cast(tokenIdx) * qkv_row + + (nq + 2 * nkv + niq) * kHeadDim; + norm_w = ik_norm_w; + } else { + return; + } + + // Store destination. Q and index_q are gathered into dedicated contiguous + // output buffers (when provided) so the downstream SM100 sparse kernel's + // flat TMA descriptor can address them as [tokens*heads, head_dim]; this + // folds the de-interleaving into the store the kernel already does, instead + // of a separate q.contiguous() copy. Everything else stays in place. + scalar_t* store_ptr = row_ptr; + if (isQ && q_out != nullptr) { + store_ptr = q_out + static_cast(tokenIdx) * nq * kHeadDim + + slot * kHeadDim; + } else if (isIQ && index_q_out != nullptr) { + // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store. + if constexpr (!kFp8Idx) { + store_ptr = index_q_out + + static_cast(tokenIdx) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim; + } + } + + // PDL: wait for the predecessor kernel (the qkv-projection GEMM that + // produces ``qkv``) to finish 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 + + // ── Load -> norm+rope (fp32) -> store back in place. ─────────────────── + float elems[kElemsPerLane]; + loadElems(row_ptr + dim_base, elems); + + if (!isV) { + int64_t const pos = positions[tokenIdx]; + scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim; + normAndRope(elems, laneId, eps, norm_w, do_rope, rotary_dim, + cos_ptr, /*apply_norm=*/norm_w != nullptr); + if constexpr (kFp8Idx) { + // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t. + if (isIQ && index_q_out != nullptr) { + storeElemsFp8(index_q_out + + static_cast(tokenIdx) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim + dim_base, + elems); + } else { + storeElems(store_ptr + dim_base, elems); + } + } else { + storeElems(store_ptr + dim_base, elems); + } + } + + // ── Cache inserts (sparse serving only). ─────────────────────────────── + if constexpr (kInsertKV) { + // Guard (not early-return) so every thread reaches the PDL trigger below. + int64_t sm = -1; + if (isK || isV) { + sm = slot_mapping[tokenIdx]; + } else if constexpr (kProcessIndex) { + if (isIK) sm = index_slot_mapping[tokenIdx]; + } + if (sm >= 0) { // skip padded / unscheduled tokens + if (isIK) { + if constexpr (kFp8Idx) { + storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems); + } else { + storeElems(index_cache + sm * kHeadDim + dim_base, elems); + } + } else if (isK || isV) { + // kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim]. + // Paging is logical (block = sm/block_size, token = sm%block_size); + // the physical NHD/HND layout is honoured via the passed strides. + int64_t const b = sm / block_size; + int64_t const t = sm % block_size; + int const kv = isK ? 0 : 1; + int64_t const off = b * kv_s_block + head * kv_s_head + + t * kv_s_token + + (kv * kHeadDim + dim_base) * kv_s_dim; + storeCacheElems(kv_cache + off, elems); + } + } + } + + // PDL: signal that this kernel is done so a dependent successor may launch + // early. No-op when PDL is not enabled on the launch. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Launch wrapper +// ──────────────────────────────────────────────────────────────────────────── +template +void launchFusedMiniMaxM3( + scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w, + scalar_t const* k_norm_w, scalar_t const* iq_norm_w, + scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache, + int64_t const* positions, int64_t const* slot_mapping, + int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache, + float const eps, int const rotary_dim, int const num_tokens, int const nq, + int const nkv, int const niq, int const block_size, + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim, bool const has_index, bool const insert_kv, + bool const process_index, bool const fp8_idx, cudaStream_t stream) { + // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the + // void* pointers per instantiation in the LAUNCH macro. + // Slot count must match the kernel's compile-time gating. + int const v_slots = insert_kv ? nkv : 0; + int const idx_slots = process_index ? niq + 1 : 0; + int const slots_per_token = nq + nkv + v_slots + idx_slots; + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int64_t const total_warps = + static_cast(num_tokens) * slots_per_token; + int const grid = + static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (grid == 0) return; + +#ifndef USE_ROCM + // 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 via cudaLaunchKernelEx. + static int const sm_version = getSMVersion(); + 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; + + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ + cudaLaunchKernelEx( \ + &config, \ + fusedMiniMaxM3QNormRopeKVInsertKernel, \ + qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, k_norm_w, \ + iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \ + index_slot_mapping, kv_cache, reinterpret_cast(index_cache), \ + eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ + kv_s_head, kv_s_token, kv_s_dim) +#else + // ROCm: standard kernel launch syntax (no PDL/stream serialization). + // clang-format off + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ + fusedMiniMaxM3QNormRopeKVInsertKernel< \ + scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \ + FP8><<>>( \ + qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, \ + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + slot_mapping, index_slot_mapping, kv_cache, \ + reinterpret_cast(index_cache), eps, rotary_dim, num_tokens, \ + nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token, \ + kv_s_dim) + // clang-format on +#endif + + if (has_index) { + if (!process_index) { + if (insert_kv) { + LAUNCH(true, true, false, false, scalar_t); + } else { + LAUNCH(true, false, false, false, scalar_t); + } + } else if (insert_kv) { + if (fp8_idx) { + LAUNCH(true, true, true, true, + uint8_t); // sparse serving, fp8 index outputs + } else { + LAUNCH(true, true, true, false, scalar_t); // sparse serving, bf16 + } + } else { + if (fp8_idx) { + LAUNCH(true, false, true, true, + uint8_t); // sparse profiling, fp8 index_q + } else { + LAUNCH(true, false, true, false, scalar_t); // sparse profiling, bf16 + } + } + } else { + // Dense layer: never has an index branch and never inserts here (the + // generic Attention layer owns the KV insert). + LAUNCH(false, false, false, false, scalar_t); + } +#undef LAUNCH +} + +} // namespace minimax_m3_fused_ops +} // namespace vllm + +// clang-format off +#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \ + vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3( \ + reinterpret_cast(qkv.data_ptr()), \ + q_out.has_value() ? reinterpret_cast(q_out->data_ptr()) : nullptr, \ + index_q_out.has_value() \ + ? reinterpret_cast(index_q_out->data_ptr()) \ + : nullptr, \ + reinterpret_cast(q_norm_weight.data_ptr()), \ + reinterpret_cast(k_norm_weight.data_ptr()), \ + process_index \ + ? reinterpret_cast(index_q_norm_weight->data_ptr()) \ + : nullptr, \ + process_index \ + ? reinterpret_cast(index_k_norm_weight->data_ptr()) \ + : nullptr, \ + reinterpret_cast(cos_sin_cache.data_ptr()), \ + reinterpret_cast(positions.data_ptr()), \ + insert_kv ? reinterpret_cast(slot_mapping->data_ptr()) \ + : nullptr, \ + (insert_kv && process_index) \ + ? reinterpret_cast( \ + effective_index_slot_mapping->data_ptr()) \ + : nullptr, \ + insert_kv ? reinterpret_cast(kv_cache->data_ptr()) : nullptr, \ + (insert_kv && process_index) \ + ? reinterpret_cast(index_cache->data_ptr()) \ + : nullptr, \ + static_cast(eps), static_cast(rotary_dim), num_tokens, nq, \ + nkv, niq, static_cast(block_size), kv_s_block, kv_s_head, \ + kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx, \ + stream) +// clang-format on + +// ──────────────────────────────────────────────────────────────────────────── +// Torch op wrapper +// ──────────────────────────────────────────────────────────────────────────── +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, // [N, qkv_row] (packs index if sparse) + torch::stable::Tensor const& q_norm_weight, // [128] + torch::stable::Tensor const& k_norm_weight, // [128] + torch::stable::Tensor const& cos_sin_cache, // [max_pos, rotary_dim] + torch::stable::Tensor const& positions, // [N] i64 + int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, // [128] + std::optional index_k_norm_weight, // [128] + int64_t num_index_heads, // niq; 0 => dense + std::optional slot_mapping, // [N] i64 + std::optional index_slot_mapping, // [N] i64 + std::optional kv_cache, // [nb,nkv,bs,2*128] + std::optional index_cache, // [nb,bs,128] + int64_t block_size, + std::optional q_out, // [N, nq*128] contiguous + std::optional + index_q_out, // [N, niq*128] contiguous + const std::string& kv_cache_dtype, bool skip_index_branch) { + STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(), + "qkv must be contiguous CUDA"); + STD_TORCH_CHECK( + qkv.scalar_type() == torch::headeronly::ScalarType::Half || + qkv.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "qkv must be float16 or bfloat16"); + STD_TORCH_CHECK( + positions.is_cuda() && + positions.scalar_type() == torch::headeronly::ScalarType::Long, + "positions must be int64 CUDA"); + STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(), + "cos_sin_cache must be contiguous CUDA"); + STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(), + "cos_sin_cache dtype must match qkv"); + STD_TORCH_CHECK( + cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim, + "cos_sin_cache shape [max_pos, rotary_dim]"); + + STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() && + k_norm_weight.scalar_type() == qkv.scalar_type(), + "q/k norm weight dtype must match qkv"); + STD_TORCH_CHECK( + q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim && + k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim, + "q/k norm weight must have 128 elements"); + STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 && + rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim, + "rotary_dim must be a positive multiple of 8 and <= 128"); + + int const num_tokens = static_cast(qkv.size(0)); + int const nq = static_cast(num_heads); + int const nkv = static_cast(num_kv_heads); + int const niq = static_cast(num_index_heads); + + // The sparse layer packs the index branch ([index_q (niq heads) | index_k + // (1 head)]) right after [q|k|v] in the same row; the dense layer does not. + bool const has_index = niq > 0; + bool const insert_kv = kv_cache.has_value(); + bool const process_index = has_index && !skip_index_branch; + vllm::Fp8KVCacheDataType const kv_dt = + vllm::get_fp8_kv_cache_data_type(kv_cache_dtype); + int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim; + int const expected_row = + (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim; + STD_TORCH_CHECK(qkv.size(1) == expected_row, + "qkv last dim must be (num_heads + 2*num_kv_heads" + " + num_index_heads + 1) * 128 for sparse, " + "(num_heads + 2*num_kv_heads) * 128 for dense"); + + // Only the sparse layer inserts here (dense lets the generic Attention layer + // own the KV write); there is no dense+insert kernel instantiation. + STD_TORCH_CHECK( + !insert_kv || has_index, + "insert mode (kv_cache) requires the index branch (sparse layer)"); + STD_TORCH_CHECK(has_index || !skip_index_branch, + "skip_index_branch requires sparse qkv rows"); + if (process_index) { + STD_TORCH_CHECK( + index_q_norm_weight.has_value() && index_k_norm_weight.has_value(), + "index branch requires both index norm weights"); + STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() && + index_k_norm_weight->scalar_type() == qkv.scalar_type(), + "index norm weights dtype must match qkv"); + STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim && + index_k_norm_weight->numel() == kHeadDim, + "index norm weights must have 128 elements"); + } + // kv_cache strides (logical shape [nb, nkv, bs, 2*head_dim]). Read straight + // off the tensor so the kernel honours whatever physical layout the attention + // backend allocated (NHD: stride order (0,2,1,3); HND: (0,1,2,3)). No new + // op argument is needed -- the strides ride along with the tensor itself. + int64_t kv_s_block = 0, kv_s_head = 0, kv_s_token = 0, kv_s_dim = 0; + torch::stable::Tensor const* effective_index_slot_mapping = nullptr; + if (insert_kv) { + STD_TORCH_CHECK( + slot_mapping.has_value() && slot_mapping->is_cuda() && + slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long, + "insert mode requires int64 CUDA slot_mapping"); + if (process_index) { + STD_TORCH_CHECK( + !index_slot_mapping.has_value() || + (index_slot_mapping->is_cuda() && + index_slot_mapping->scalar_type() == + torch::headeronly::ScalarType::Long && + index_slot_mapping->numel() == slot_mapping->numel()), + "index_slot_mapping must be int64 CUDA with slot_mapping length"); + } + // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage. + if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) { + STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(), + "auto kv_cache dtype must match qkv"); + } else { + STD_TORCH_CHECK( + kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte, + "fp8 kv_cache must use uint8 storage"); + } + // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3. + if (process_index) { + STD_TORCH_CHECK( + index_cache.has_value() && + (index_cache->scalar_type() == qkv.scalar_type() || + index_cache->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn), + "insert mode requires index_cache matching qkv dtype or fp8 e4m3"); + } + STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1, + "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous " + "content dim (stride(3)==1)"); + kv_s_block = kv_cache->stride(0); + kv_s_head = kv_cache->stride(1); + kv_s_token = kv_cache->stride(2); + kv_s_dim = kv_cache->stride(3); + if (process_index) { + effective_index_slot_mapping = index_slot_mapping.has_value() + ? &index_slot_mapping.value() + : &slot_mapping.value(); + } + } + // Optional contiguous gather targets: when given, the normed/roped q (and + // index_q) are written here instead of in place, so callers avoid a separate + // .contiguous() copy. index_q_out only makes sense on the sparse path. + if (q_out.has_value()) { + STD_TORCH_CHECK( + q_out->is_cuda() && q_out->is_contiguous() && + q_out->scalar_type() == qkv.scalar_type(), + "q_out must be a contiguous CUDA tensor matching qkv dtype"); + STD_TORCH_CHECK( + q_out->numel() == static_cast(num_tokens) * nq * kHeadDim, + "q_out must have num_tokens * num_heads * 128 elements"); + } + if (index_q_out.has_value()) { + STD_TORCH_CHECK(process_index, + "index_q_out requires index branch processing"); + STD_TORCH_CHECK( + index_q_out->is_cuda() && index_q_out->is_contiguous() && + (index_q_out->scalar_type() == qkv.scalar_type() || + index_q_out->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn), + "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3"); + STD_TORCH_CHECK(index_q_out->numel() == + static_cast(num_tokens) * niq * kHeadDim, + "index_q_out must have num_tokens * num_index_heads * 128 " + "elements"); + } + + // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while + // q/k/v + q_out stay qkv dtype. Both index outputs must agree. + auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn; + bool const fp8_idx = + process_index && + ((index_cache.has_value() && index_cache->scalar_type() == kFp8) || + (index_q_out.has_value() && index_q_out->scalar_type() == kFp8)); + if (fp8_idx) { + STD_TORCH_CHECK( + !index_cache.has_value() || index_cache->scalar_type() == kFp8, + "fp8 index path: index_cache must be fp8 e4m3"); + STD_TORCH_CHECK( + !index_q_out.has_value() || index_q_out->scalar_type() == kFp8, + "fp8 index path: index_q_out must be fp8 e4m3"); + } + + const torch::stable::accelerator::DeviceGuard device_guard( + qkv.get_device_index()); + auto stream = get_current_cuda_stream(qkv.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] { + using st = scalar_t; + DISPATCH_BY_KV_CACHE_DTYPE(qkv.scalar_type(), kv_cache_dtype, + CALL_FUSED_MINIMAX_M3); + }); +} + +#undef CALL_FUSED_MINIMAX_M3 diff --git a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu b/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu index c9b7ee9e4e9..a8e6d32a1be 100644 --- a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu +++ b/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu @@ -22,7 +22,7 @@ #include "async_util.cuh" #include "../cuda_compat.h" -#include "../type_convert.cuh" +#include "type_convert.cuh" #include "dispatch_utils.h" #define CHECK_TYPE(x, st) \ diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 37df6be329f..878b44df936 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -2,16 +2,16 @@ #include "torch_utils.h" -#include "../cub_helpers.h" +#include "cub_helpers.h" #include "../core/batch_invariant.hpp" -#include "../type_convert.cuh" +#include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" namespace vllm { // TODO(woosuk): Further optimize this kernel. -template +template __global__ void rms_norm_kernel( scalar_t* __restrict__ out, // [..., hidden_size] const scalar_t* __restrict__ input, // [..., hidden_size] @@ -20,20 +20,26 @@ __global__ void rms_norm_kernel( const int64_t input_stride_d4, // input.stride(-4) const int64_t input_shape_d2, // input.size(-2) const int64_t input_shape_d3, // input.size(-3) - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size] or + // [num_groups, hidden_size]; + // null if !HasWeight + const int64_t weight_stride, // 0 or weight.stride(0) const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; const scalar_t* input_row; + const scalar_t* weight_row; if constexpr (NUM_DIMS == 2) { // 2D for layernorm normal case [batch_size, hidden] input_row = input + blockIdx.x * input_stride_d2; + weight_row = weight + blockIdx.x * weight_stride; } else if constexpr (NUM_DIMS == 3) { // 3D for q/k norm [batch_size, num_heads, head_size] int batch_idx = blockIdx.x / input_shape_d2; int head_idx = blockIdx.x % input_shape_d2; input_row = input + batch_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } else if constexpr (NUM_DIMS == 4) { // 4D for transformers model_impl qk norm [batch, seq, head, head_dim] int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2); @@ -42,6 +48,7 @@ __global__ void rms_norm_kernel( int head_idx = remaining % input_shape_d2; input_row = input + batch_idx * input_stride_d4 + seq_idx * input_stride_d3 + head_idx * input_stride_d2; + weight_row = weight + batch_idx * weight_stride; } auto vec_op = [&variance](const vec_n_t& vec) { @@ -69,16 +76,24 @@ __global__ void rms_norm_kernel( scalar_t* out_row = out + blockIdx.x * hidden_size; auto* v_in = reinterpret_cast*>(input_row); - auto* v_w = reinterpret_cast*>(weight); + auto* v_w = reinterpret_cast*>(weight_row); auto* v_out = reinterpret_cast*>(out_row); for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { vec_n_t dst; vec_n_t src1 = v_in[i]; - vec_n_t src2 = v_w[i]; + vec_n_t src2; + if constexpr (HasWeight) { + src2 = v_w[i]; + } #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - dst.val[j] = static_cast(x * s_variance) * src2.val[j]; + if constexpr (HasWeight) { + float w = static_cast(src2.val[j]); + dst.val[j] = static_cast(x * s_variance * w); + } else { + dst.val[j] = static_cast(x * s_variance); + } } v_out[i] = dst; } @@ -88,13 +103,13 @@ __global__ void rms_norm_kernel( Additional optimizations we can make in this case are packed and vectorized operations, which help with the memory latency bottleneck. */ -template +template __global__ std::enable_if_t<(width > 0) && _typeConvert::exists> fused_add_rms_norm_kernel( scalar_t* __restrict__ input, // [..., hidden_size] const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight const float epsilon, const int num_tokens, const int hidden_size) { // Sanity checks on our vector struct and type-punned pointer arithmetic static_assert(std::is_pod_v<_f16Vec>); @@ -136,13 +151,22 @@ fused_add_rms_norm_kernel( int id = blockIdx.x * vec_hidden_size + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec res = residual_v[id]; - _f16Vec w = weight_v[idx]; _f16Vec out; using Converter = _typeConvert; + if constexpr (HasWeight) { + _f16Vec w = weight_v[idx]; #pragma unroll - for (int j = 0; j < width; ++j) { - float x = Converter::convert(res.data[j]); - out.data[j] = Converter::convert(x * s_variance) * w.data[j]; + 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); + } + } else { +#pragma unroll + for (int j = 0; j < width; ++j) { + float x = Converter::convert(res.data[j]); + out.data[j] = Converter::convert(x * s_variance); + } } input_v[strided_id] = out; } @@ -151,13 +175,13 @@ fused_add_rms_norm_kernel( /* Generic fused_add_rms_norm_kernel The width field is not used here but necessary for other specializations. */ -template +template __global__ std::enable_if_t<(width == 0) || !_typeConvert::exists> fused_add_rms_norm_kernel( scalar_t* __restrict__ input, // [..., hidden_size] const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] - const scalar_t* __restrict__ weight, // [hidden_size] + const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight const float epsilon, const int num_tokens, const int hidden_size) { __shared__ float s_variance; float variance = 0.0f; @@ -181,23 +205,38 @@ 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]; + if constexpr (HasWeight) { + float w = (float)weight[idx]; + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); + } else { + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance); + } } } } // namespace vllm -void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] - torch::stable::Tensor& input, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] - double epsilon) { +void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] + torch::stable::Tensor& input, // [..., hidden_size] + std::optional weight, double epsilon) { STD_TORCH_CHECK(out.is_contiguous()); if (input.stride(-1) != 1) { input = torch::stable::contiguous(input); } STD_TORCH_CHECK(input.stride(-1) == 1); - STD_TORCH_CHECK(weight.is_contiguous()); + int64_t weight_stride = 0; + if (weight.has_value()) { + STD_TORCH_CHECK(weight->is_contiguous()); + if (weight->dim() == 1) { + STD_TORCH_CHECK(weight->size(0) == input.size(-1)); + } else if (weight->dim() == 2) { + STD_TORCH_CHECK(weight->size(0) == input.size(0)); + STD_TORCH_CHECK(weight->size(-1) == input.size(-1)); + weight_stride = weight->stride(0); + } else { + STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D"); + } + } int hidden_size = input.size(-1); @@ -215,46 +254,69 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); + const bool has_weight = weight.has_value(); VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] { VLLM_STABLE_DISPATCH_FLOATING_TYPES( input.scalar_type(), "rms_norm_kernel", [&] { + const scalar_t* weight_ptr = + has_weight ? weight->const_data_ptr() : nullptr; const int calculated_vec_size = std::gcd(16 / sizeof(scalar_t), hidden_size); const int block_size = std::min(hidden_size / calculated_vec_size, max_block_size); dim3 block(block_size); VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] { - vllm::rms_norm_kernel - <<>>( - out.mutable_data_ptr(), - input.const_data_ptr(), input_stride_d2, - input_stride_d3, input_stride_d4, input_shape_d2, - input_shape_d3, weight.const_data_ptr(), epsilon, - num_tokens, hidden_size); + if (has_weight) { + vllm::rms_norm_kernel + <<>>( + out.mutable_data_ptr(), + input.const_data_ptr(), input_stride_d2, + input_stride_d3, input_stride_d4, input_shape_d2, + input_shape_d3, weight_ptr, weight_stride, epsilon, + num_tokens, hidden_size); + } else { + vllm::rms_norm_kernel + <<>>( + out.mutable_data_ptr(), + input.const_data_ptr(), input_stride_d2, + input_stride_d3, input_stride_d4, input_shape_d2, + input_shape_d3, weight_ptr, /*weight_stride=*/0, epsilon, + num_tokens, hidden_size); + } }); }); }); } -#define LAUNCH_FUSED_ADD_RMS_NORM(width) \ - VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ - input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ - vllm::fused_add_rms_norm_kernel \ - <<>>( \ - input.mutable_data_ptr(), input_stride, \ - residual.mutable_data_ptr(), \ - weight.const_data_ptr(), epsilon, num_tokens, \ - hidden_size); \ +#define LAUNCH_FUSED_ADD_RMS_NORM(width, has_weight) \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ + if (has_weight) { \ + vllm::fused_add_rms_norm_kernel \ + <<>>( \ + input.mutable_data_ptr(), input_stride, \ + residual.mutable_data_ptr(), \ + weight->const_data_ptr(), epsilon, num_tokens, \ + hidden_size); \ + } else { \ + vllm::fused_add_rms_norm_kernel \ + <<>>( \ + input.mutable_data_ptr(), input_stride, \ + residual.mutable_data_ptr(), nullptr, epsilon, \ + num_tokens, hidden_size); \ + } \ }); void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] torch::stable::Tensor& residual, // [..., hidden_size] - torch::stable::Tensor& weight, // [hidden_size] + std::optional weight, double epsilon) { - STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type()); STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type()); STD_TORCH_CHECK(residual.is_contiguous()); - STD_TORCH_CHECK(weight.is_contiguous()); + if (weight.has_value()) { + STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type()); + STD_TORCH_CHECK(weight->is_contiguous()); + } int hidden_size = input.size(-1); int64_t input_stride = input.stride(-2); int num_tokens = input.numel() / hidden_size; @@ -269,30 +331,33 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); - /*If the tensor types are FP16/BF16, try to use the optimized kernel - with packed + vectorized ops. - Max optimization is achieved with a width-8 vector of FP16/BF16s - since we can load at most 128 bits at once in a global memory op. - However, this requires each tensor's data to be aligned to 16 - bytes. - */ + constexpr int vector_width = 8; + constexpr int req_alignment_bytes = vector_width * 2; auto inp_ptr = reinterpret_cast(input.data_ptr()); auto res_ptr = reinterpret_cast(residual.data_ptr()); - auto wt_ptr = reinterpret_cast(weight.data_ptr()); - constexpr int vector_width = 8; - constexpr int req_alignment_bytes = - vector_width * 2; // vector_width * sizeof(bfloat16 or float16) (float32 - // falls back to non-vectorized version anyway) - bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && - res_ptr % req_alignment_bytes == 0 && - wt_ptr % req_alignment_bytes == 0; bool offsets_are_multiple_of_vector_width = hidden_size % vector_width == 0 && input_stride % vector_width == 0; bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); - if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && - !batch_invariant_launch) { - LAUNCH_FUSED_ADD_RMS_NORM(8); + const bool has_weight = weight.has_value(); + if (has_weight) { + auto wt_ptr = reinterpret_cast(weight->data_ptr()); + bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && + res_ptr % req_alignment_bytes == 0 && + wt_ptr % req_alignment_bytes == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && + !batch_invariant_launch) { + LAUNCH_FUSED_ADD_RMS_NORM(8, true); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0, true); + } } else { - LAUNCH_FUSED_ADD_RMS_NORM(0); + bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && + res_ptr % req_alignment_bytes == 0; + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && + !batch_invariant_launch) { + LAUNCH_FUSED_ADD_RMS_NORM(8, false); + } else { + LAUNCH_FUSED_ADD_RMS_NORM(0, false); + } } } diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index 32f3495f4e9..f3bf8882e77 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -9,10 +9,10 @@ #include "torch_utils.h" -#include "../cub_helpers.h" +#include "cub_helpers.h" #include "../core/batch_invariant.hpp" #include "../quantization/w8a8/fp8/common.cuh" -#include "../type_convert.cuh" +#include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" @@ -66,8 +66,13 @@ __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]); - // Multiply in weight's native dtype to match rms_norm_kernel. - scalar_t out_norm = static_cast(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(static_cast(out_norm), scale_inv); @@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel( #pragma unroll for (int i = 0; i < width; ++i) { float x = Converter::convert(res.data[i]); - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - HipT out_norm_h = Converter::convert(x * s_variance) * w.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); } @@ -183,8 +192,10 @@ 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]; - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - scalar_t out_norm = static_cast(x * s_variance) * weight[idx]; + float w = (float)weight[idx]; + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. + scalar_t out_norm = static_cast(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/minimax_reduce_rms_kernel.cu b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu index d9af0f5efe0..58d61b353d6 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu @@ -249,7 +249,7 @@ __global__ void __launch_bounds__(1024) LamportComm comm(params.workspace, params.rank); int clear_access = comm.clear_size / kElemsPerAccess; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif for (int idx = access_id; idx < tot_access; idx += access_stride, token_id += token_stride) { @@ -313,7 +313,7 @@ __global__ void __launch_bounds__(1024) } comm.update(params.size_q * NRanks); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -384,7 +384,7 @@ __global__ void __launch_bounds__(1024) DType norm_weight[kElemsPerAccess]{}; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif if (is_q) { if (is_valid_q) { @@ -596,7 +596,7 @@ __global__ void __launch_bounds__(1024) } } // end group loop #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif int clear_access = static_cast(comm.clear_size / kElemsPerAccess); @@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { } // namespace tensorrt_llm } // namespace vllm -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps) { - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - 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 = const_cast(input.const_data_ptr()); - allreduce_params.rms_gamma = const_cast(norm_weight.const_data_ptr()); - allreduce_params.rms_eps = static_cast(eps); - allreduce_params.stream = get_current_cuda_stream(input.get_device_index()); - - torch::stable::Tensor rms_norm_out = torch::stable::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::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, diff --git a/csrc/minimax_reduce_rms_kernel.h b/csrc/libtorch_stable/minimax_reduce_rms_kernel.h similarity index 100% rename from csrc/minimax_reduce_rms_kernel.h rename to csrc/libtorch_stable/minimax_reduce_rms_kernel.h diff --git a/csrc/moe/dsv3_router_gemm_bf16_out.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu similarity index 81% rename from csrc/moe/dsv3_router_gemm_bf16_out.cu rename to csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu index b11ba991b26..8695d1e8084 100644 --- a/csrc/moe/dsv3_router_gemm_bf16_out.cu +++ b/csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu @@ -18,14 +18,11 @@ * limitations under the License. */ -#include -#include +#include #include #include -#include "dsv3_router_gemm_utils.h" - // Custom FMA implementation using PTX assembly instructions __device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, float2 const& c) { @@ -81,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output( } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // Process the GEMM in chunks @@ -166,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output( } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -289,3 +286,52 @@ template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>( template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>( __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144) +template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>( + __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu new file mode 100644 index 00000000000..c5ebaf11f10 --- /dev/null +++ b/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu @@ -0,0 +1,210 @@ +/* + * Adapted from SGLang's sgl-kernel implementation, which was adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu + * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp + * + * Copyright (c) 2019-2023, 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 "libtorch_stable/torch_utils.h" + +#include +#include + +#include + +namespace { + +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} + +} // namespace + +static constexpr int DEFAULT_NUM_EXPERTS = 256; +static constexpr int KIMI_K2_NUM_EXPERTS = 384; +static constexpr int DEFAULT_HIDDEN_DIM = 7168; +static constexpr int GLM_5_HIDDEN_DIM = 6144; + +template +void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, + cudaStream_t stream); + +template +void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, + T const* mat_b, cudaStream_t stream); + +template +struct LoopUnroller { + static void unroll_float_output(int num_tokens, float* output, + __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, + cudaStream_t stream) { + if (num_tokens == kBegin) { + invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts, + kHiddenDim>(output, input, weights, stream); + } else { + LoopUnroller::unroll_float_output(num_tokens, output, input, + weights, stream); + } + } + + static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, + __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, + cudaStream_t stream) { + if (num_tokens == kBegin) { + invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts, + kHiddenDim>(output, input, weights, stream); + } else { + LoopUnroller::unroll_bf16_output(num_tokens, output, input, + weights, stream); + } + } +}; + +template +struct LoopUnroller { + static void unroll_float_output(int num_tokens, float* output, + __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, + cudaStream_t stream) { + if (num_tokens == kEnd) { + invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( + output, input, weights, stream); + } else { + throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); + } + } + + static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, + __nv_bfloat16 const* input, + __nv_bfloat16 const* weights, + cudaStream_t stream) { + if (num_tokens == kEnd) { + invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( + output, input, weights, stream); + } else { + throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); + } + } +}; + +void dsv3_router_gemm( + torch::stable::Tensor& output, // [num_tokens, num_experts] + torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim] + torch::stable::Tensor const& mat_b // [num_experts, hidden_dim] +) { + STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); + + const int num_tokens = mat_a.size(0); + const int num_experts = mat_b.size(0); + const int hidden_dim = mat_a.size(1); + + STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1), + "mat_a and mat_b must have the same hidden_dim"); + STD_TORCH_CHECK( + hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM, + "Expected hidden_dim=", DEFAULT_HIDDEN_DIM, + " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim); + STD_TORCH_CHECK( + num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS, + "Expected num_experts=", DEFAULT_NUM_EXPERTS, + " or num_experts=", KIMI_K2_NUM_EXPERTS, + ", but got num_experts=", num_experts); + // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim. + STD_TORCH_CHECK( + hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS, + "hidden_dim=", GLM_5_HIDDEN_DIM, + " only supports num_experts=", DEFAULT_NUM_EXPERTS, + ", but got num_experts=", num_experts); + STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, + "currently num_tokens must be less than or equal to 16 for " + "router_gemm"); + STD_TORCH_CHECK( + mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "mat_a must be bf16"); + STD_TORCH_CHECK( + mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "mat_b must be bf16"); + STD_TORCH_CHECK( + output.scalar_type() == torch::headeronly::ScalarType::Float || + output.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "output must be float32 or bf16"); + + const torch::stable::accelerator::DeviceGuard device_guard( + mat_a.get_device_index()); + const int sm = getSMVersion(); + STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90"); + + const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index()); + + __nv_bfloat16 const* a_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); + __nv_bfloat16 const* b_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); + + if (output.scalar_type() == torch::headeronly::ScalarType::Float) { + float* out_ptr = reinterpret_cast(output.mutable_data_ptr()); + if (hidden_dim == DEFAULT_HIDDEN_DIM) { + if (num_experts == DEFAULT_NUM_EXPERTS) { + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, + DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens, + out_ptr, a_ptr, + b_ptr, stream); + } else { + LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, + DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens, + out_ptr, a_ptr, + b_ptr, stream); + } + } else { // GLM_5_HIDDEN_DIM + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, + GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr, + a_ptr, b_ptr, stream); + } + } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + __nv_bfloat16* out_ptr = + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + if (hidden_dim == DEFAULT_HIDDEN_DIM) { + if (num_experts == DEFAULT_NUM_EXPERTS) { + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, + DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens, + out_ptr, a_ptr, + b_ptr, stream); + } else { + LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, + DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens, + out_ptr, a_ptr, + b_ptr, stream); + } + } else { // GLM_5_HIDDEN_DIM + LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, + GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr, + a_ptr, b_ptr, stream); + } + } +} + +STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { + m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm)); +} diff --git a/csrc/moe/dsv3_router_gemm_float_out.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu similarity index 81% rename from csrc/moe/dsv3_router_gemm_float_out.cu rename to csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu index 2756cba0b14..58a2b44ae2f 100644 --- a/csrc/moe/dsv3_router_gemm_float_out.cu +++ b/csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu @@ -18,14 +18,11 @@ * limitations under the License. */ -#include -#include +#include #include #include -#include "dsv3_router_gemm_utils.h" - // Custom FMA implementation using PTX assembly instructions __device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, float2 const& c) { @@ -81,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output( } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // Process the GEMM in chunks @@ -166,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output( } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -289,3 +286,52 @@ template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>( template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144) +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/csrc/moe/grouped_topk_kernels.cu b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu similarity index 92% rename from csrc/moe/grouped_topk_kernels.cu rename to csrc/libtorch_stable/moe/grouped_topk_kernels.cu index 6a4dad3be7c..da9ef44d03f 100644 --- a/csrc/moe/grouped_topk_kernels.cu +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu @@ -18,9 +18,14 @@ * limitations under the License. */ #include "moeTopKFuncs.cuh" -#include -#include + +#include +#include + +#include "libtorch_stable/torch_utils.h" + #include +#include #include #include #include @@ -43,7 +48,8 @@ static constexpr int NumTopGroupScores = 2; static constexpr int DefaultMaxNumTopExperts = 8; static constexpr int MaxSupportedTopExperts = 22; static constexpr int MaxNumTopGroups = 4; - +// The empirical value for small batch +static constexpr int PDLEnableTokens = 16; namespace warp_topk { template @@ -559,8 +565,8 @@ __global__ void grouped_topk_fused_kernel( T* s_group_scores = reinterpret_cast(ptr_u); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); // I think all prolog can be put before - // acqbulk because it's ptr arithmetic + cudaGridDependencySynchronize(); // I think all prolog can be put before + // acqbulk because it's ptr arithmetic #endif // phase 1: per-group scan @@ -604,7 +610,7 @@ __global__ void grouped_topk_fused_kernel( topk_values[i] = 1.0f / static_cast(topk_i32); } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif return; } @@ -665,7 +671,7 @@ __global__ void grouped_topk_fused_kernel( } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -890,7 +896,8 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, int64_t const num_experts, int64_t const n_group, int64_t const topk_group, int64_t const topk, bool const renormalize, double const routed_scaling_factor, - bool enable_pdl = false, cudaStream_t const stream = 0) { + const bool enable_pdl = false, + cudaStream_t const stream = 0) { cudaLaunchConfig_t config; config.stream = stream; cudaLaunchAttribute attrs[1]; @@ -978,7 +985,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, int64_t const num_tokens, int64_t const num_experts, \ int64_t const n_group, int64_t const topk_group, int64_t const topk, \ bool const renormalize, double const routed_scaling_factor, \ - bool enable_pdl, cudaStream_t const stream); + const bool enable_pdl, cudaStream_t const stream); INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID); INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID); @@ -1001,38 +1008,43 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE); } // end namespace moe } // namespace vllm -std::tuple grouped_topk( - torch::Tensor const& scores, int64_t n_group, int64_t topk_group, +std::tuple grouped_topk( + torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize, double routed_scaling_factor, - torch::Tensor const& bias, int64_t scoring_func = 0) { - auto data_type = scores.scalar_type(); - auto bias_type = bias.scalar_type(); - auto input_size = scores.sizes(); - int64_t num_tokens = input_size[0]; - int64_t num_experts = input_size[1]; - TORCH_CHECK(input_size.size() == 2, "scores must be a 2D Tensor"); - TORCH_CHECK(n_group > 0, "n_group must be positive"); - TORCH_CHECK(topk > 0, "topk must be positive"); - TORCH_CHECK(topk_group > 0, "topk_group must be positive"); - TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group"); - TORCH_CHECK(num_experts % n_group == 0, - "num_experts should be divisible by n_group"); - TORCH_CHECK(n_group <= 32, - "n_group should be smaller than or equal to 32 for now"); - TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now"); - TORCH_CHECK(topk <= topk_group * (num_experts / n_group), - "topk must be <= topk_group * (num_experts / n_group)"); - TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE || - scoring_func == vllm::moe::SCORING_SIGMOID, - "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)"); + torch::stable::Tensor const& bias, int64_t scoring_func = 0) { + const auto data_type = scores.scalar_type(); + const auto bias_type = bias.scalar_type(); + STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor"); + const int64_t num_tokens = scores.size(0); + const int64_t num_experts = scores.size(1); + STD_TORCH_CHECK(n_group > 0, "n_group must be positive"); + STD_TORCH_CHECK(topk > 0, "topk must be positive"); + STD_TORCH_CHECK(topk_group > 0, "topk_group must be positive"); + STD_TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group"); + STD_TORCH_CHECK(num_experts % n_group == 0, + "num_experts should be divisible by n_group"); + STD_TORCH_CHECK(n_group <= 32, + "n_group should be smaller than or equal to 32 for now"); + STD_TORCH_CHECK(topk <= 32, + "topk should be smaller than or equal to 32 for now"); + STD_TORCH_CHECK(topk <= topk_group * (num_experts / n_group), + "topk must be <= topk_group * (num_experts / n_group)"); + STD_TORCH_CHECK( + scoring_func == vllm::moe::SCORING_NONE || + scoring_func == vllm::moe::SCORING_SIGMOID, + "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)"); // Always output float32 for topk_values (eliminates Python-side conversion) - torch::Tensor topk_values = torch::empty( - {num_tokens, topk}, torch::dtype(torch::kFloat32).device(torch::kCUDA)); - torch::Tensor topk_indices = torch::empty( - {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA)); + auto topk_values = torch::stable::new_empty( + scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float); + auto topk_indices = torch::stable::new_empty( + scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int); + const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens; - auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device()); + const torch::stable::accelerator::DeviceGuard device_guard( + scores.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(scores.get_device_index()); auto const sf = static_cast(scoring_func); #define LAUNCH_KERNEL_SF(T, BiasT, IdxT) \ @@ -1045,7 +1057,7 @@ std::tuple grouped_topk( reinterpret_cast(topk_indices.mutable_data_ptr()), \ reinterpret_cast(bias.data_ptr()), num_tokens, \ num_experts, n_group, topk_group, topk, renormalize, \ - routed_scaling_factor, false, stream); \ + routed_scaling_factor, pdl_flag, stream); \ break; \ case vllm::moe::SCORING_SIGMOID: \ vllm::moe::invokeNoAuxTc( \ @@ -1054,10 +1066,10 @@ std::tuple grouped_topk( reinterpret_cast(topk_indices.mutable_data_ptr()), \ reinterpret_cast(bias.data_ptr()), num_tokens, \ num_experts, n_group, topk_group, topk, renormalize, \ - routed_scaling_factor, false, stream); \ + routed_scaling_factor, pdl_flag, stream); \ break; \ default: \ - throw std::invalid_argument("Unsupported scoring_func"); \ + STD_TORCH_CHECK(false, "Unsupported scoring_func"); \ break; \ } \ } while (0) @@ -1065,17 +1077,18 @@ std::tuple grouped_topk( #define LAUNCH_KERNEL(T, IdxT) \ do { \ switch (bias_type) { \ - case torch::kFloat16: \ + case torch::headeronly::ScalarType::Half: \ LAUNCH_KERNEL_SF(T, half, IdxT); \ break; \ - case torch::kFloat32: \ + case torch::headeronly::ScalarType::Float: \ LAUNCH_KERNEL_SF(T, float, IdxT); \ break; \ - case torch::kBFloat16: \ + case torch::headeronly::ScalarType::BFloat16: \ LAUNCH_KERNEL_SF(T, __nv_bfloat16, IdxT); \ break; \ default: \ - throw std::invalid_argument( \ + STD_TORCH_CHECK( \ + false, \ "Invalid bias dtype, only supports float16, float32, and " \ "bfloat16"); \ break; \ @@ -1083,22 +1096,22 @@ std::tuple grouped_topk( } while (0) switch (data_type) { - case torch::kFloat16: + case torch::headeronly::ScalarType::Half: // Handle Float16 LAUNCH_KERNEL(half, int32_t); break; - case torch::kFloat32: + case torch::headeronly::ScalarType::Float: // Handle Float32 LAUNCH_KERNEL(float, int32_t); break; - case torch::kBFloat16: + case torch::headeronly::ScalarType::BFloat16: // Handle BFloat16 LAUNCH_KERNEL(__nv_bfloat16, int32_t); break; default: // Handle other data types - throw std::invalid_argument( - "Invalid dtype, only supports float16, float32, and bfloat16"); + STD_TORCH_CHECK( + false, "Invalid dtype, only supports float16, float32, and bfloat16"); break; } #undef LAUNCH_KERNEL diff --git a/csrc/moe/marlin_moe_wna16/.gitignore b/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore similarity index 100% rename from csrc/moe/marlin_moe_wna16/.gitignore rename to csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore diff --git a/csrc/moe/marlin_moe_wna16/generate_kernels.py b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py similarity index 99% rename from csrc/moe/marlin_moe_wna16/generate_kernels.py rename to csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py index 6ddda1d51db..64b47b607bb 100644 --- a/csrc/moe/marlin_moe_wna16/generate_kernels.py +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py @@ -302,7 +302,7 @@ def generate_new_kernels(): if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: kernel_selector_str += ( "else if (a_type == vllm::kFE4M3fn)\n" - " TORCH_CHECK(false, " + " STD_TORCH_CHECK(false, " '"marlin kernel with fp8 activation is not built.");' ) diff --git a/csrc/moe/marlin_moe_wna16/kernel.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h similarity index 95% rename from csrc/moe/marlin_moe_wna16/kernel.h rename to csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h index 09ed1a470bd..783736ab509 100644 --- a/csrc/moe/marlin_moe_wna16/kernel.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/moe/marlin_moe_wna16/marlin_template.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h similarity index 99% rename from csrc/moe/marlin_moe_wna16/marlin_template.h rename to csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h index 9858df94573..04f90101be4 100644 --- a/csrc/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" -#include "quantization/marlin/dequant.h" -#include "quantization/marlin/marlin_mma.h" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/dequant.h" +#include "libtorch_stable/quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/moe/marlin_moe_wna16/ops.cu b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu similarity index 62% rename from csrc/moe/marlin_moe_wna16/ops.cu rename to csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu index 82cba2978b1..177eefa2c6f 100644 --- a/csrc/moe/marlin_moe_wna16/ops.cu +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu @@ -24,7 +24,15 @@ #endif #include "kernel.h" -#include "core/registration.h" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ @@ -350,18 +358,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, bool m_block_size_8 = moe_block_size == 8; bool is_a_8bit = a_type.size_bits() == 8; - TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, - ", ", prob_n, ", ", prob_k, "]"); + STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); int group_blocks = 0; if (has_act_order) { if (is_k_full) { - TORCH_CHECK(group_size != -1); + STD_TORCH_CHECK(group_size != -1); group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } else { - TORCH_CHECK(group_size == 0); + STD_TORCH_CHECK(group_size == 0); group_blocks = 0; } } else { @@ -369,8 +377,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, group_blocks = -1; } else { group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } } @@ -407,7 +415,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, else if (moe_block_size == 64) kernel = permute_cols_kernel<64>; else - TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size); + STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size); // avoid ">>>" being formatted to "> > >" // clang-format off @@ -428,25 +436,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + STD_TORCH_CHECK(max_shared_mem > 0); int major_capability, minor_capability; cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, dev); cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, dev); - TORCH_CHECK(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); int stages = 4; if (major_capability == 7 && minor_capability == 5) { stages = 2; - TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); + STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); } if (a_type == vllm::kFE4M3fn) { - TORCH_CHECK(major_capability * 10 + minor_capability >= 89, - "FP8 only support Ada Lovelace or newer GPUs."); - TORCH_CHECK( + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); + STD_TORCH_CHECK( major_capability * 10 + minor_capability == 89 || major_capability == 12, "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " @@ -460,10 +468,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64}; if (blocks_per_sm == -1) blocks_per_sm = 1; exec_cfg = exec_config_t{blocks_per_sm, thread_tfg}; - TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); - TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); + STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); } else { // Auto config exec_cfg = determine_exec_config( @@ -484,19 +492,19 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int thread_k_blocks = thread_k / 16; int thread_n_blocks = thread_n / 16; - TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, - prob_m, prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, - is_a_8bit, stages, max_shared_mem), - "Invalid thread config: thread_m_blocks = ", thread_m_blocks, - ", thread_k = ", thread_tfg.thread_k, - ", thread_n = ", thread_tfg.thread_n, - ", num_threads = ", thread_tfg.num_threads, " for MKN = [", - prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, - ", group_size = ", group_size, - ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, - ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, - ", max_shared_mem = ", max_shared_mem); + STD_TORCH_CHECK( + is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages, + max_shared_mem), + "Invalid thread config: thread_m_blocks = ", thread_m_blocks, + ", thread_k = ", thread_tfg.thread_k, + ", thread_n = ", thread_tfg.thread_n, + ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", + prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, ", has_act_order = ", has_act_order, + ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, + ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem); int sh_cache_size = get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, @@ -509,13 +517,13 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, num_threads, is_zp_float, stages); if (kernel == MarlinDefault) { - TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, - ", ", prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, - ", num_bits = ", num_bits); + STD_TORCH_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits); } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -532,75 +540,81 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, } // namespace MARLIN_NAMESPACE_NAME -torch::Tensor moe_wna16_marlin_gemm( - torch::Tensor& a, std::optional c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& a_scales_or_none, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::Tensor& workspace, - torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids, - torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights, - int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, - vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, - int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, - bool is_zp_float, int64_t thread_k, int64_t thread_n, +torch::stable::Tensor moe_wna16_marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional const& perm_or_none, + torch::stable::Tensor& workspace, torch::stable::Tensor& sorted_token_ids, + torch::stable::Tensor& expert_ids, + torch::stable::Tensor& num_tokens_past_padded, + torch::stable::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k, + bool mul_topk_weights, vllm::ScalarTypeId const& b_type_id, int64_t size_m, + int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, int64_t thread_n, int64_t blocks_per_sm) { vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; - auto c_dtype = a.dtype(); - if (a.scalar_type() == at::ScalarType::Half) { + auto c_dtype = a.scalar_type(); + if (a.scalar_type() == torch::headeronly::ScalarType::Half) { a_type_id = vllm::kFloat16.id(); c_type_id = vllm::kFloat16.id(); - } else if (a.scalar_type() == at::ScalarType::BFloat16) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { a_type_id = vllm::kBFloat16.id(); c_type_id = vllm::kBFloat16.id(); } else { - c_dtype = b_scales.dtype(); - if (b_scales.scalar_type() == at::ScalarType::Half) { + c_dtype = b_scales.scalar_type(); + if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { c_type_id = vllm::kBFloat16.id(); - TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); - torch::Tensor c = c_or_none.value(); - c_dtype = c.dtype(); + STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); + torch::stable::Tensor c = c_or_none.value(); + c_dtype = c.scalar_type(); - if (c.scalar_type() == at::ScalarType::Half) { + if (c.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (c.scalar_type() == at::ScalarType::BFloat16) { + } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { - TORCH_CHECK(false, "unsupported c dtype"); + STD_TORCH_CHECK(false, "unsupported c dtype"); } } - if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) { a_type_id = vllm::kFE4M3fn.id(); - } else if (a.scalar_type() == at::ScalarType::Char) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) { a_type_id = vllm::kS8.id(); } else { - TORCH_CHECK(false, "unsupported `a` scalar_type"); + STD_TORCH_CHECK(false, "unsupported `a` scalar_type"); } } s_type_id = c_type_id; if (b_type_id == vllm::kFE2M1f.id()) { - if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn) { s_type_id = vllm::kFE4M3fn.id(); - } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } else { - TORCH_CHECK(false, - "When b_type = float4_e2m1f, b_scale scalar type must be", - "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); + STD_TORCH_CHECK( + false, "When b_type = float4_e2m1f, b_scale scalar type must be", + "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); } } else if (b_type_id == vllm::kFE4M3fn.id() && - b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } @@ -613,58 +627,60 @@ torch::Tensor moe_wna16_marlin_gemm( int num_experts = b_q_weight.size(0); if (moe_block_size != 8) { - TORCH_CHECK(moe_block_size % 16 == 0, - "unsupported moe_block_size=", moe_block_size); - TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64, - "unsupported moe_block_size=", moe_block_size); + STD_TORCH_CHECK(moe_block_size % 16 == 0, + "unsupported moe_block_size=", moe_block_size); + STD_TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64, + "unsupported moe_block_size=", moe_block_size); } // Verify A - TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), - ", size_m = ", size_m); - TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), - ", size_k = ", size_k); + STD_TORCH_CHECK(a.size(0) == size_m, + "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + STD_TORCH_CHECK(a.size(1) == size_k, + "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); // Verify B - TORCH_CHECK( + STD_TORCH_CHECK( size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k, " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1), - "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), - ", size_k = ", size_k, - ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK( + STD_TORCH_CHECK( + (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1), + "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), + ", size_k = ", size_k, + ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); + STD_TORCH_CHECK( b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0, "b_q_weight.size(2) = ", b_q_weight.size(2), " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); int actual_size_n = (b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor; - TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, - ", actual_size_n = ", actual_size_n); + STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); // Verify device and strides - TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); - TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); + STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous"); - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); - TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); - torch::Tensor a_scales; - auto options = torch::TensorOptions().dtype(c_dtype).device(a.device()); - auto options_fp32 = - torch::TensorOptions().dtype(at::kFloat).device(a.device()); + torch::stable::Tensor a_scales; + constexpr auto kFloat = torch::headeronly::ScalarType::Float; if (a_scales_or_none.has_value()) { a_scales = a_scales_or_none.value(); - TORCH_CHECK(a_type.size_bits() == 8, - "a_scales can only be used for 8bit activation."); + STD_TORCH_CHECK(a_type.size_bits() == 8, + "a_scales can only be used for 8bit activation."); } else { - a_scales = torch::empty({0}, options_fp32); - TORCH_CHECK(a_type.size_bits() != 8, - "the a_scales parameter must be passed for 8bit activation."); + a_scales = torch::stable::new_empty(a, {0}, kFloat); + STD_TORCH_CHECK( + a_type.size_bits() != 8, + "the a_scales parameter must be passed for 8bit activation."); } // sms: number of SMs to use for the kernel @@ -672,82 +688,84 @@ torch::Tensor moe_wna16_marlin_gemm( cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device()); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - torch::Tensor c; + torch::stable::accelerator::DeviceGuard device_guard(a.get_device_index()); + torch::stable::Tensor c; if (c_or_none.has_value()) { c = c_or_none.value(); - TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); - TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); - TORCH_CHECK(c.size(0) == size_m * top_k, - "Shape mismatch: c.size(0) = ", c.size(0), - ", size_m * topk = ", size_m * top_k); - TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1), - ", size_n = ", size_n); + STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); + STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); + STD_TORCH_CHECK(c.size(0) == size_m * top_k, + "Shape mismatch: c.size(0) = ", c.size(0), + ", size_m * topk = ", size_m * top_k); + STD_TORCH_CHECK(c.size(1) == size_n, + "Shape mismatch: c.size(1) = ", c.size(1), + ", size_n = ", size_n); } else { - c = torch::empty({size_m * top_k, size_n}, options); + c = torch::stable::new_empty(a, {size_m * top_k, size_n}, c_dtype); } // Alloc C tmp buffer that is going to be used for the global reduce - torch::Tensor c_tmp; + torch::stable::Tensor c_tmp; if (use_fp32_reduce && !use_atomic_add) { // max num of threadblocks is sms * 4 long max_c_tmp_size = min( (long)size_n * sorted_token_ids.size(0), (long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n); if (moe_block_size == 8) max_c_tmp_size *= 2; - c_tmp = torch::empty({max_c_tmp_size}, options_fp32); + c_tmp = torch::stable::new_empty(a, {max_c_tmp_size}, kFloat); } else { - c_tmp = torch::empty({0}, options_fp32); + c_tmp = torch::stable::new_empty(a, {0}, kFloat); } // Detect groupsize and act_order int num_groups = -1; int group_size = -1; - int rank = b_scales.sizes().size(); - TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3"); - TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2), - " is not size_n = ", size_n); + int rank = b_scales.dim(); + STD_TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3"); + STD_TORCH_CHECK(b_scales.size(2) == size_n, + "b_scales dim 2 = ", b_scales.size(2), + " is not size_n = ", size_n); num_groups = b_scales.size(1); - torch::Tensor g_idx, perm, a_tmp; + torch::stable::Tensor g_idx, perm, a_tmp; if (g_idx_or_none.has_value() && perm_or_none.has_value()) { g_idx = g_idx_or_none.value(); perm = perm_or_none.value(); - TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); // Verify g_idx and perm - TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || - (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", g_idx.size(-1), - " and perm.size(-1) = ", perm.size(-1), - ", where size_k = ", size_k); + STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || + (g_idx.size(-1) == size_k && perm.size(-1) == size_k), + "Unexpected g_idx.size(-1) = ", g_idx.size(-1), + " and perm.size(-1) = ", perm.size(-1), + ", where size_k = ", size_k); } else { - g_idx = torch::empty({0}, options); - perm = torch::empty({0}, options); - a_tmp = torch::empty({0}, options); + g_idx = torch::stable::new_empty(a, {0}, c_dtype); + perm = torch::stable::new_empty(a, {0}, c_dtype); + a_tmp = torch::stable::new_empty(a, {0}, c_dtype); } bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; if (has_act_order) { - a_tmp = torch::empty({size_m * top_k, size_k}, options); + a_tmp = torch::stable::new_empty(a, {size_m * top_k, size_k}, c_dtype); if (is_k_full) { - TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); + STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); group_size = size_k / num_groups; } else { group_size = 0; } } else { - a_tmp = torch::empty({0}, options); + a_tmp = torch::stable::new_empty(a, {0}, c_dtype); if (num_groups > 1) { - TORCH_CHECK( + STD_TORCH_CHECK( size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by b_scales.size(1) = ", b_scales.size(1)); group_size = size_k / num_groups; @@ -756,119 +774,125 @@ torch::Tensor moe_wna16_marlin_gemm( } } - torch::Tensor global_scale; + torch::stable::Tensor global_scale; if (global_scale_or_none.has_value()) { global_scale = global_scale_or_none.value(); - TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "global_scale can only be used for nvfp4 format."); + STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, + "global_scale can only be used for nvfp4 format."); } else { - global_scale = torch::empty({0}, options_fp32); - TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), - "the global_scale parameter must be passed for nvfp4 format."); + global_scale = torch::stable::new_empty(a, {0}, kFloat); + STD_TORCH_CHECK( + !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), + "the global_scale parameter must be passed for nvfp4 format."); } bool has_bias = b_bias_or_none.has_value(); - torch::Tensor b_bias; + torch::stable::Tensor b_bias; if (has_bias) { b_bias = b_bias_or_none.value(); - TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); - TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); - TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n"); - TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1"); + STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); + STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); + STD_TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(1) != size_n"); + STD_TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1"); } else { - b_bias = torch::empty({0}, options); + b_bias = torch::stable::new_empty(a, {0}, c_dtype); } - torch::Tensor b_zeros; + torch::stable::Tensor b_zeros; if (b_zeros_or_none.has_value()) { b_zeros = b_zeros_or_none.value(); - TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); - TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); + STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); + STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); } else { - b_zeros = torch::empty({0}, options); + b_zeros = torch::stable::new_empty(a, {0}, c_dtype); } bool has_zp = b_zeros.size(-1) > 0; if (has_zp) { - TORCH_CHECK( + STD_TORCH_CHECK( b_type == vllm::kU4 || b_type == vllm::kU8, "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str()); } else { - TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || - b_type == vllm::kS4 || b_type == vllm::kS8 || - b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, - "b_type must be uint4b8, uint8b128, int4, int8, " - "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", - b_type.str()); + STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || + b_type == vllm::kS4 || b_type == vllm::kS8 || + b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, + "b_type must be uint4b8, uint8b128, int4, int8, " + "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", + b_type.str()); } if (has_zp && is_zp_float) { - TORCH_CHECK(a.scalar_type() == at::ScalarType::Half, - "Computation type must be float16 (half) when using float zero " - "points."); + STD_TORCH_CHECK( + a.scalar_type() == torch::headeronly::ScalarType::Half, + "Computation type must be float16 (half) when using float zero " + "points."); } // Verify b_zeros if (has_zp) { - int rank = b_zeros.sizes().size(); - TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3"); + int rank = b_zeros.dim(); + STD_TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3"); if (is_zp_float) { - TORCH_CHECK(b_zeros.size(2) == size_n, - "b_zeros dim 2 = ", b_zeros.size(2), - " is not size_n = ", size_n); - TORCH_CHECK(num_groups == b_zeros.size(1), - "b_zeros dim 1 = ", b_zeros.size(1), - " is not num_groups = ", num_groups); - TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); + STD_TORCH_CHECK(b_zeros.size(2) == size_n, + "b_zeros dim 2 = ", b_zeros.size(2), + " is not size_n = ", size_n); + STD_TORCH_CHECK(num_groups == b_zeros.size(1), + "b_zeros dim 1 = ", b_zeros.size(1), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); } else { - TORCH_CHECK(b_zeros.size(1) == num_groups, - "b_zeros dim 1 = ", b_zeros.size(1), - " is not num_groups = ", num_groups); - TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor, - "b_zeros dim 2 = ", b_zeros.size(2), - " is not size_n / pack_factor = ", size_n / pack_factor); + STD_TORCH_CHECK(b_zeros.size(1) == num_groups, + "b_zeros dim 1 = ", b_zeros.size(1), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor, + "b_zeros dim 2 = ", b_zeros.size(2), + " is not size_n / pack_factor = ", size_n / pack_factor); } } // Verify workspace size - TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, - "size_n = ", size_n, ", is not divisible by min_thread_n = ", - MARLIN_NAMESPACE_NAME::min_thread_n); + STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, + "size_n = ", size_n, ", is not divisible by min_thread_n = ", + MARLIN_NAMESPACE_NAME::min_thread_n); int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n; int min_workspace_size = min( max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4); - TORCH_CHECK(workspace.numel() >= min_workspace_size, - "workspace.numel = ", workspace.numel(), - " is below min_workspace_size = ", min_workspace_size); + STD_TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); int dev = a.get_device(); - TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float, - "scalar type of a_scales must be float"); - TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, - "scalar type of global_scale must be float"); + STD_TORCH_CHECK( + a_scales.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of a_scales must be float"); + STD_TORCH_CHECK( + global_scale.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of global_scale must be float"); if (a_type.size_bits() == 16) { - TORCH_CHECK( + STD_TORCH_CHECK( a.scalar_type() == c.scalar_type(), "scalar type of a must be the same with c for 16 bit activation"); } MARLIN_NAMESPACE_NAME::marlin_mm( - a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), - b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(), - global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), - perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(), - topk_weights.data_ptr(), moe_block_size, num_experts, top_k, - mul_topk_weights, size_m, size_n, size_k, workspace.data_ptr(), a_type, - b_type, c_type, s_type, has_bias, has_act_order, is_k_full, has_zp, - num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev), + a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(), + c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), + a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), + global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), + g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), + a_tmp.mutable_data_ptr(), sorted_token_ids.mutable_data_ptr(), + expert_ids.mutable_data_ptr(), num_tokens_past_padded.mutable_data_ptr(), + topk_weights.mutable_data_ptr(), moe_block_size, num_experts, top_k, + mul_topk_weights, size_m, size_n, size_k, workspace.mutable_data_ptr(), + a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full, + has_zp, num_groups, group_size, dev, get_current_cuda_stream(dev), thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce, is_zp_float); return c; } -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm); +STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { + m.impl("moe_wna16_marlin_gemm", TORCH_BOX(&moe_wna16_marlin_gemm)); } diff --git a/csrc/moe/moeTopKFuncs.cuh b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh similarity index 100% rename from csrc/moe/moeTopKFuncs.cuh rename to csrc/libtorch_stable/moe/moeTopKFuncs.cuh diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu similarity index 58% rename from csrc/moe/moe_align_sum_kernels.cu rename to csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index a8fa59b1939..152bc116edb 100644 --- a/csrc/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -1,14 +1,18 @@ -#include -#include -#include +#include #include -#include -#include +#include +#include +#include +#include +#include +#include -#include "../cuda_compat.h" -#include "../dispatch_utils.h" -#include "core/math.hpp" +#include "../../cuda_compat.h" +#include "libtorch_stable/core/math.hpp" +#include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/quantization/vectorization.cuh" +#include "libtorch_stable/torch_utils.h" #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) @@ -78,6 +82,21 @@ __global__ void batched_moe_align_block_size_kernel( } } // namespace batched_moe_align_block_size +template +__device__ __forceinline__ int get_local_expert_id( + size_t idx, const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ expert_map, int32_t num_experts, + bool has_expert_map) { + int expert_id = topk_ids[idx]; + if (expert_id >= num_experts || expert_id < 0) { + return -1; + } + if (has_expert_map) { + expert_id = expert_map[expert_id]; + } + return expert_id; +} + template __device__ void _moe_align_block_size( const scalar_t* __restrict__ topk_ids, @@ -122,20 +141,15 @@ __device__ void _moe_align_block_size( const size_t stride = blockDim.x; for (size_t i = tid; i < numel; i += stride) { - int expert_id = topk_ids[i]; - if (expert_id >= num_experts) { - continue; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int warp_idx = expert_id / experts_per_warp; + int expert_offset = expert_id % experts_per_warp; + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], + mask); } - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid experts - if (expert_id == -1) continue; - } - int warp_idx = expert_id / experts_per_warp; - int expert_offset = expert_id % experts_per_warp; - int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; - atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], - mask); } __syncthreads(); @@ -223,14 +237,12 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid expert - if (expert_id == -1) continue; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; } - int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; - tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; } __syncthreads(); @@ -272,18 +284,16 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid expert - if (expert_id == -1) continue; - } - int32_t rank_post_pad = - tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + int32_t rank_post_pad = + tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; - if (token_mask == nullptr || token_mask[i / topk_num]) { - sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; - ++tokens_cnts[tid * num_experts + expert_id]; + if (token_mask == nullptr || token_mask[i / topk_num]) { + sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; + ++tokens_cnts[tid * num_experts + expert_id]; + } } } } @@ -299,22 +309,15 @@ __device__ void _count_and_sort_expert_tokens( const size_t stride = blockDim.x * gridDim.y; for (size_t i = tid; i < numel; i += stride) { - int32_t expert_id = topk_ids[i]; - if (expert_id >= num_experts) { - continue; - } - - if (has_expert_map) { - expert_id = expert_map[expert_id]; - // filter invalid experts - if (expert_id == -1) continue; - } - - if (token_mask == nullptr || token_mask[i / topk_num]) { - int32_t rank_post_pad = atomicAdd( - &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); - sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = - i; + if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, + num_experts, has_expert_map); + expert_id != -1) { + if (token_mask == nullptr || token_mask[i / topk_num]) { + int32_t rank_post_pad = atomicAdd( + &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); + sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = + i; + } } } } @@ -346,19 +349,102 @@ __global__ void count_and_sort_expert_tokens_kernel( max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map); } +// Reduce the topk expert outputs per token (summed in fp32). The output is +// dense [num_tokens, d]; the input is addressed by its strides so non- +// contiguous inputs work without a copy. A 16B-vectorized path is used when +// the hidden dim is contiguous (innermost stride 1) and aligned; otherwise a +// scalar kernel reads via arbitrary strides. topk is a compile-time constant +// for common values and runtime otherwise. + +// Elements per 16-byte vector (8 for bf16/fp16, 4 for fp32). +template +constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t); + template -__global__ void moe_sum_kernel( - scalar_t* __restrict__ out, // [..., d] - const scalar_t* __restrict__ input, // [..., topk, d] - const int d) { - const int64_t token_idx = blockIdx.x; - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { - scalar_t x = 0.0; +__global__ void moe_sum_vec_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous + const int64_t num_tokens, const int d, const int64_t stride_token, + const int64_t stride_topk) { + using vec_t = vllm::vec_n_t>; // 16-byte pack + constexpr int VEC = MOE_SUM_VEC; + const int64_t n_vec = d / VEC; + const int64_t total = num_tokens * n_vec; + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; + i += (int64_t)gridDim.x * blockDim.x) { + const int64_t token = i / n_vec; + const int64_t v = i % n_vec; + const scalar_t* in_tok = input + token * stride_token + v * VEC; + + float acc[VEC]; +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] = 0.f; + #pragma unroll for (int k = 0; k < TOPK; ++k) { - x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]); + vec_t packed = *reinterpret_cast(in_tok + k * stride_topk); +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] += static_cast(packed.val[j]); } - out[token_idx * d + idx] = x; + + vec_t outp; +#pragma unroll + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast(acc[j]); + *reinterpret_cast(out + token * d + v * VEC) = outp; + } +} + +// Runtime-topk variant of the above. +template +__global__ void moe_sum_vec_dynamic_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous + const int64_t num_tokens, const int d, const int topk, + const int64_t stride_token, const int64_t stride_topk) { + using vec_t = vllm::vec_n_t>; + constexpr int VEC = MOE_SUM_VEC; + const int64_t n_vec = d / VEC; + const int64_t total = num_tokens * n_vec; + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total; + i += (int64_t)gridDim.x * blockDim.x) { + const int64_t token = i / n_vec; + const int64_t v = i % n_vec; + const scalar_t* in_tok = input + token * stride_token + v * VEC; + + float acc[VEC]; +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] = 0.f; + + for (int k = 0; k < topk; ++k) { + vec_t packed = *reinterpret_cast(in_tok + k * stride_topk); +#pragma unroll + for (int j = 0; j < VEC; ++j) acc[j] += static_cast(packed.val[j]); + } + + vec_t outp; +#pragma unroll + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast(acc[j]); + *reinterpret_cast(out + token * d + v * VEC) = outp; + } +} + +// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims +// (including a non-contiguous hidden stride) via per-element strided reads. +template +__global__ void moe_sum_scalar_kernel( + scalar_t* __restrict__ out, // [num_tokens, d], contiguous + const scalar_t* __restrict__ input, // [num_tokens, topk, d] + const int d, const int topk, const int64_t stride_token, + const int64_t stride_topk, const int64_t stride_hidden) { + const int64_t token_idx = blockIdx.x; + const scalar_t* in_tok = input + token_idx * stride_token; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + float x = 0.f; + for (int k = 0; k < topk; ++k) { + x += static_cast( + VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden])); + } + out[token_idx * d + idx] = static_cast(x); } } @@ -492,12 +578,15 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel( // taken from // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc -void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, - int64_t block_size, torch::Tensor sorted_token_ids, - torch::Tensor experts_ids, - torch::Tensor num_tokens_post_pad, - std::optional maybe_expert_map) { - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); +void moe_align_block_size( + torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map) { + const torch::stable::accelerator::DeviceGuard device_guard( + topk_ids.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(topk_ids.get_device_index()); int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; @@ -506,19 +595,18 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; // BlockScan uses 1024 threads and assigns one thread per expert. - TORCH_CHECK(padded_num_experts < 1024, - "padded_num_experts must be less than 1024"); - auto options_int = - torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device()); + STD_TORCH_CHECK(padded_num_experts < 1024, + "padded_num_experts must be less than 1024"); bool has_expert_map = maybe_expert_map.has_value(); - torch::Tensor expert_map; + torch::stable::Tensor expert_map; if (has_expert_map) { expert_map = maybe_expert_map.value(); } else { - expert_map = torch::empty({0}, options_int); + expert_map = torch::stable::new_empty(topk_ids, {0}, + torch::headeronly::ScalarType::Int); } - VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES( + VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES( topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] { // calc needed amount of shared mem for `cumsum` tensors bool small_batch_expert_mode = @@ -538,16 +626,17 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, scalar_t, fill_threads>; small_batch_expert_kernel<<<1, fill_threads + threads, shared_mem_size, stream>>>( - topk_ids.data_ptr(), - sorted_token_ids.data_ptr(), - experts_ids.data_ptr(), - num_tokens_post_pad.data_ptr(), - expert_map.data_ptr(), num_experts, block_size, - topk_ids.numel(), sorted_token_ids.size(0), topk_ids.size(1), - has_expert_map); + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, block_size, topk_ids.numel(), + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); } else { - torch::Tensor cumsum_buffer = - torch::empty({num_experts + 1}, options_int); + torch::stable::Tensor cumsum_buffer = torch::stable::new_empty( + topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); auto align_kernel = vllm::moe::moe_align_block_size_kernel; size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp); @@ -558,14 +647,16 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, // blockIdx.x == 0: counting experts and aligning // blockIdx.x == 1: filling sorted_token_ids align_kernel<<<2, threads, shared_mem_size, stream>>>( - topk_ids.data_ptr(), - sorted_token_ids.data_ptr(), - experts_ids.data_ptr(), - num_tokens_post_pad.data_ptr(), - expert_map.data_ptr(), num_experts, padded_num_experts, - experts_per_warp, block_size, topk_ids.numel(), - cumsum_buffer.data_ptr(), sorted_token_ids.size(0), - topk_ids.size(1), has_expert_map); + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, padded_num_experts, experts_per_warp, block_size, + topk_ids.numel(), + reinterpret_cast(cumsum_buffer.mutable_data_ptr()), + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); const int block_threads = std::min(256, (int)threads); const int num_blocks = @@ -577,9 +668,10 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, auto sort_kernel = vllm::moe::count_and_sort_expert_tokens_kernel; sort_kernel<<>>( - topk_ids.data_ptr(), - sorted_token_ids.data_ptr(), - cumsum_buffer.data_ptr(), expert_map.data_ptr(), + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum_buffer.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), topk_ids.numel(), num_experts, sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); } @@ -588,111 +680,156 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts, void batched_moe_align_block_size(int64_t max_tokens_per_batch, int64_t block_size, - torch::Tensor const& batch_num_tokens, - torch::Tensor sorted_ids, - torch::Tensor batch_ids, - torch::Tensor num_tokens_post_pad) { + const torch::stable::Tensor& batch_num_tokens, + torch::stable::Tensor sorted_ids, + torch::stable::Tensor batch_ids, + torch::stable::Tensor num_tokens_post_pad) { namespace batched_kernel = vllm::moe::batched_moe_align_block_size; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + batch_num_tokens.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(batch_num_tokens.get_device_index()); int32_t const B = batch_num_tokens.size(0); int32_t const num_blocks_per_batch = round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size; int32_t const num_blocks = num_blocks_per_batch * B; int64_t const sorted_ids_size = num_blocks * block_size; - TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size); - TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size); - TORCH_CHECK(num_tokens_post_pad.size(0) == 1); - TORCH_CHECK(B <= batched_kernel::num_threads); + STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size); + STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size); + STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1); + STD_TORCH_CHECK(B <= batched_kernel::num_threads); batched_kernel::batched_moe_align_block_size_kernel<<< batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>( - B, max_tokens_per_batch, block_size, batch_num_tokens.data_ptr(), - sorted_ids.data_ptr(), batch_ids.data_ptr(), - num_tokens_post_pad.data_ptr()); + B, max_tokens_per_batch, block_size, + reinterpret_cast(batch_num_tokens.const_data_ptr()), + reinterpret_cast(sorted_ids.mutable_data_ptr()), + reinterpret_cast(batch_ids.mutable_data_ptr()), + reinterpret_cast(num_tokens_post_pad.mutable_data_ptr())); } -void moe_sum(torch::Tensor& input, // [num_tokens, topk, hidden_size] - torch::Tensor& output) // [num_tokens, hidden_size] +void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size] + torch::stable::Tensor& output) // [num_tokens, hidden_size] { + // Output is dense and written in place, so it must be contiguous. The input + // is read by its strides (no copy); only the hidden dim needs to be + // contiguous to take the vectorized path. + STD_TORCH_CHECK(output.is_contiguous(), + "moe_sum expects a contiguous output"); + const int hidden_size = input.size(-1); - const auto num_tokens = output.numel() / hidden_size; + const int64_t num_tokens = output.numel() / hidden_size; const int topk = input.size(1); + const int64_t stride_token = input.stride(0); + const int64_t stride_topk = input.stride(1); + const int64_t stride_hidden = input.stride(2); - dim3 grid(num_tokens); - dim3 block(std::min(hidden_size, 1024)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(output)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + output.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(output.get_device_index()); - switch (topk) { - case 2: - VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - output.data_ptr(), input.data_ptr(), - hidden_size); - }); - break; +#define LAUNCH_MOE_SUM_VEC(TOPK) \ + vllm::moe::moe_sum_vec_kernel \ + <<>>( \ + out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk) - case 3: - VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - output.data_ptr(), input.data_ptr(), - hidden_size); - }); - break; + VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] { + constexpr int VEC = vllm::moe::MOE_SUM_VEC; + constexpr int WIDTH = VEC * sizeof(scalar_t); // 16 bytes + auto* out_ptr = reinterpret_cast(output.mutable_data_ptr()); + auto* in_ptr = reinterpret_cast(input.const_data_ptr()); - case 4: - VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] { - vllm::moe::moe_sum_kernel<<>>( - output.data_ptr(), input.data_ptr(), - hidden_size); - }); - break; - - default: - at::sum_out(output, input, 1); - break; - } + // Vectorize along hidden only when it is contiguous (innermost stride 1), + // a whole number of vectors, and every row offset stays 16B-aligned. + const bool can_vec = (stride_hidden == 1) && (hidden_size % VEC == 0) && + (stride_token % VEC == 0) && + (stride_topk % VEC == 0) && + (reinterpret_cast(in_ptr) % WIDTH == 0) && + (reinterpret_cast(out_ptr) % WIDTH == 0); + if (can_vec) { + const int64_t n_vec = hidden_size / VEC; + const int64_t total = num_tokens * n_vec; + const int block = 256; + const dim3 grid(std::min((total + block - 1) / block, 65535)); + switch (topk) { + case 1: + LAUNCH_MOE_SUM_VEC(1); + break; + case 2: + LAUNCH_MOE_SUM_VEC(2); + break; + case 4: + LAUNCH_MOE_SUM_VEC(4); + break; + case 6: + LAUNCH_MOE_SUM_VEC(6); + break; + case 8: + LAUNCH_MOE_SUM_VEC(8); + break; + case 9: + LAUNCH_MOE_SUM_VEC(9); + break; + default: + vllm::moe::moe_sum_vec_dynamic_kernel + <<>>(out_ptr, in_ptr, num_tokens, + hidden_size, topk, + stride_token, stride_topk); + break; + } + } else { + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + vllm::moe::moe_sum_scalar_kernel<<>>( + out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk, + stride_hidden); + } + }); +#undef LAUNCH_MOE_SUM_VEC } void moe_lora_align_block_size( - torch::Tensor topk_ids, torch::Tensor token_lora_mapping, + torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping, int64_t num_experts, int64_t block_size, int64_t max_loras, int64_t max_num_tokens_padded, int64_t max_num_m_blocks, - torch::Tensor sorted_token_ids, torch::Tensor expert_ids, - torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled, - torch::Tensor lora_ids, std::optional maybe_expert_map) { + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids, + std::optional maybe_expert_map) { const int topk_num = topk_ids.size(1); - TORCH_CHECK(block_size > 0, "block_size should be greater than 0. "); + STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. "); int device_max_shared_mem; - auto dev = topk_ids.get_device(); + int dev = topk_ids.get_device_index(); + const torch::stable::accelerator::DeviceGuard device_guard(dev); cudaDeviceGetAttribute(&device_max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(dev); int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; // BlockScan uses 1024 threads and assigns one thread per expert. - TORCH_CHECK(padded_num_experts < 1024, - "padded_num_experts must be less than 1024"); + STD_TORCH_CHECK(padded_num_experts < 1024, + "padded_num_experts must be less than 1024"); - auto options_int = - torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device()); - torch::Tensor token_mask = - torch::empty({max_loras * topk_ids.size(0)}, options_int); + torch::stable::Tensor token_mask = + torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)}, + torch::headeronly::ScalarType::Int); bool has_expert_map = maybe_expert_map.has_value(); - torch::Tensor expert_map; + torch::stable::Tensor expert_map; if (has_expert_map) { expert_map = maybe_expert_map.value(); } else { - expert_map = torch::empty({0}, options_int); + expert_map = torch::stable::new_empty(topk_ids, {0}, + torch::headeronly::ScalarType::Int); } - VLLM_DISPATCH_INTEGRAL_TYPES( + VLLM_STABLE_DISPATCH_INTEGRAL_TYPES( topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] { bool small_batch_expert_mode = (topk_ids.numel() < 1024) && (num_experts <= 64); @@ -703,7 +840,7 @@ void moe_lora_align_block_size( (num_thread + 1) * num_experts * sizeof(int32_t) + (num_experts + 1) * sizeof(int32_t); if (shared_mem > device_max_shared_mem) { - TORCH_CHECK(false, "Shared memory usage exceeds device limit."); + STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit."); } // threadIdx.x >= fill_threads: counting experts and aligning @@ -714,7 +851,7 @@ void moe_lora_align_block_size( auto kernel = vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel< scalar_t, fill_threads>; - AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( + STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( (void*)kernel, shared_mem)); // Grid size is (max_loras + 1) because active_lora_ids has length // max_loras + 1: sorted-unique values of token_lora_mapping, which @@ -725,15 +862,21 @@ void moe_lora_align_block_size( // MoE-LoRA kernels. This mirrors the fix made for the Triton // _fused_moe_lora_kernel grid in vllm-project/vllm#32277. kernel<<>>( - topk_ids.data_ptr(), - token_lora_mapping.data_ptr(), block_size, - expert_map.data_ptr(), num_experts, max_loras, - topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks, - sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), topk_num, - num_tokens_post_pad.data_ptr(), - adapter_enabled.data_ptr(), lora_ids.data_ptr(), - token_mask.data_ptr(), has_expert_map); + reinterpret_cast(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(token_mask.mutable_data_ptr()), + has_expert_map); } else { int num_thread = 1024; dim3 blockDim(num_thread); @@ -742,8 +885,9 @@ void moe_lora_align_block_size( size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t); // cumsum buffer - torch::Tensor cumsum = - torch::zeros({max_loras * (num_experts + 1)}, options_int); + torch::stable::Tensor cumsum = torch::stable::new_zeros( + topk_ids, {max_loras * (num_experts + 1)}, + torch::headeronly::ScalarType::Int); auto align_kernel = vllm::moe::moe_lora_align_block_size_kernel; @@ -759,16 +903,23 @@ void moe_lora_align_block_size( // blockIdx.x % 2 == 1: filling sorted_token_ids align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size, stream>>>( - topk_ids.data_ptr(), - token_lora_mapping.data_ptr(), block_size, - expert_map.data_ptr(), num_experts, max_loras, - topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks, - sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), topk_num, - num_tokens_post_pad.data_ptr(), - adapter_enabled.data_ptr(), cumsum.data_ptr(), - WARP_SIZE, padded_num_experts, lora_ids.data_ptr(), - token_mask.data_ptr(), has_expert_map); + reinterpret_cast(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), WARP_SIZE, + padded_num_experts, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(token_mask.mutable_data_ptr()), + has_expert_map); const int block_threads = std::min(256, (int)num_thread); const int num_blocks = @@ -785,12 +936,16 @@ void moe_lora_align_block_size( vllm::moe::lora_count_and_sort_expert_tokens_kernel; sort_kernel<<>>( - topk_ids.data_ptr(), - sorted_token_ids.data_ptr(), cumsum.data_ptr(), - expert_map.data_ptr(), topk_ids.numel(), num_experts, - max_num_tokens_padded, topk_num, token_mask.data_ptr(), - max_loras, lora_ids.data_ptr(), - adapter_enabled.data_ptr(), has_expert_map); + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num, + reinterpret_cast(token_mask.mutable_data_ptr()), + max_loras, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + has_expert_map); } }); } \ No newline at end of file diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h new file mode 100644 index 00000000000..b60d2d548f5 --- /dev/null +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -0,0 +1,88 @@ +#pragma once + +#include + +#include +#include + +void topk_softmax(torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + std::optional bias); + +void topk_sigmoid(torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + std::optional bias, + double routed_scaling_factor); + +void topk_softplus_sqrt( + torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + double routed_scaling_factor, + const std::optional& correction_bias, + const std::optional& input_ids, + const std::optional& tid2eid); + +void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output); + +void moe_align_block_size( + torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map); + +void batched_moe_align_block_size( + int64_t max_tokens_per_batch, int64_t block_size, + const torch::stable::Tensor& expert_num_tokens, + torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad); + +void moe_lora_align_block_size( + torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping, + int64_t num_experts, int64_t block_size, int64_t max_loras, + int64_t max_num_tokens_padded, int64_t max_num_m_blocks, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids, + std::optional maybe_expert_map); +#ifndef USE_ROCM +torch::stable::Tensor moe_wna16_gemm( + torch::stable::Tensor input, torch::stable::Tensor output, + torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales, + std::optional b_qzeros, + std::optional topk_weights, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, int64_t top_k, + int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K, + int64_t bit); + +std::tuple grouped_topk( + const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group, + int64_t topk, bool renormalize, double routed_scaling_factor, + const torch::stable::Tensor& bias, int64_t scoring_func); +#endif + +bool moe_permute_unpermute_supported(); + +int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, + int64_t num_expert); + +void shuffle_rows(const torch::stable::Tensor& input_tensor, + const torch::stable::Tensor& dst2src_map, + torch::stable::Tensor& output_tensor); + +#ifndef USE_ROCM +// DeepSeek V3 optimized router GEMM kernel for SM90+ +// Computes output = mat_a @ mat_b.T where: +// mat_a: [num_tokens, hidden_dim] in bf16 +// mat_b: [num_experts, hidden_dim] in bf16 +// output: [num_tokens, num_experts] in bf16 or fp32 +// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168 +void dsv3_router_gemm(torch::stable::Tensor& output, + const torch::stable::Tensor& mat_a, + const torch::stable::Tensor& mat_b); +#endif diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu new file mode 100644 index 00000000000..52cf68442f1 --- /dev/null +++ b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu @@ -0,0 +1,325 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "core/registration.h" +#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "libtorch_stable/torch_utils.h" + +#include + +// moe_permute kernels require at least CUDA 12.0 +#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) + +namespace { + +int64_t product_integers(torch::headeronly::IntHeaderOnlyArrayRef sizes) { + int64_t numel = 1; + for (int64_t s : sizes) { + numel *= s; + } + return numel; +} + +torch::stable::Tensor maybe_allocate_tensor( + const std::optional& maybe_tensor, + torch::headeronly::IntHeaderOnlyArrayRef expected_sizes, + torch::headeronly::ScalarType dtype, torch::stable::Device device, + char const* name) { + auto expected_numel = product_integers(expected_sizes); + if (maybe_tensor.has_value()) { + auto tensor = maybe_tensor.value(); + STD_TORCH_CHECK(tensor.device() == device, name, + " must be on the same device"); + STD_TORCH_CHECK(tensor.scalar_type() == dtype, name, + " has incorrect dtype"); + STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + STD_TORCH_CHECK(tensor.numel() >= expected_numel, name, + " is too small for the requested shape"); + auto flat_tensor = torch::stable::view(tensor, {tensor.numel()}); + return torch::stable::view( + torch::stable::narrow(flat_tensor, 0, 0, expected_numel), + expected_sizes); + } + return torch::stable::empty(expected_sizes, dtype, std::nullopt, device); +} + +} // namespace + +int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, + int64_t n_expert) { + return static_cast( + CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert)); +} + +void moe_permute_impl( + const torch::stable::Tensor& input, // [n_token, hidden] + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] + const std::optional& expert_map, // [n_expert] + int64_t n_expert, int64_t n_local_expert, int64_t topk, + torch::stable::Tensor& permuted_input, // [permuted_size, hidden] + torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] + torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] + torch::stable::Tensor& permuted_idx, // [permute_size] + const std::optional& maybe_sort_workspace, + const std::optional& maybe_permuted_experts_id, + const std::optional& maybe_sorted_row_idx, + const std::optional& maybe_topk_ids_for_sort) { + STD_TORCH_CHECK(expert_first_token_offset.scalar_type() == + torch::headeronly::ScalarType::Long, + "expert_first_token_offset must be int64"); + STD_TORCH_CHECK(topk_ids.scalar_type() == torch::headeronly::ScalarType::Int, + "topk_ids must be int32"); + STD_TORCH_CHECK( + token_expert_indices.scalar_type() == torch::headeronly::ScalarType::Int, + "token_expert_indices must be int32"); + STD_TORCH_CHECK( + inv_permuted_idx.scalar_type() == torch::headeronly::ScalarType::Int, + "inv_permuted_idx must be int32"); + STD_TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1, + "expert_first_token_offset shape != n_local_expert+1"); + STD_TORCH_CHECK( + inv_permuted_idx.sizes().equals(token_expert_indices.sizes()), + "token_expert_indices shape must be same as inv_permuted_idx"); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + auto device = input.device(); + auto n_token = input.sizes()[0]; + auto n_hidden = input.sizes()[1]; + auto expanded_rows = n_token * topk; + auto stream = get_current_cuda_stream(input.get_device_index()); + + auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert); + auto sort_workspace = maybe_allocate_tensor( + maybe_sort_workspace, {sorter_size}, torch::headeronly::ScalarType::Char, + device, "sort_workspace"); + auto permuted_experts_id = maybe_allocate_tensor( + maybe_permuted_experts_id, topk_ids.sizes(), + torch::headeronly::ScalarType::Int, device, "permuted_experts_id"); + auto sorted_row_idx = maybe_allocate_tensor( + maybe_sorted_row_idx, inv_permuted_idx.sizes(), + torch::headeronly::ScalarType::Int, device, "sorted_row_idx"); + + CubKeyValueSorter sorter{}; + int64_t* valid_num_ptr = nullptr; + torch::stable::Tensor topk_ids_for_sort = topk_ids; + + if (expert_map.has_value()) { + const int* expert_map_ptr = get_ptr(expert_map.value()); + valid_num_ptr = + get_ptr(expert_first_token_offset) + n_local_expert; + topk_ids_for_sort = maybe_allocate_tensor( + maybe_topk_ids_for_sort, topk_ids.sizes(), + torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort"); + torch::stable::copy_(topk_ids_for_sort, topk_ids); + preprocessTopkIdLauncher(get_ptr(topk_ids_for_sort), n_token * topk, + expert_map_ptr, n_expert, stream); + } + + sortAndScanExpert( + get_ptr(topk_ids_for_sort), get_ptr(token_expert_indices), + get_ptr(permuted_experts_id), get_ptr(sorted_row_idx), + get_ptr(expert_first_token_offset), n_token, n_expert, + n_local_expert, topk, sorter, get_ptr(sort_workspace), stream); + + MOE_DISPATCH(input.scalar_type(), [&] { + expandInputRowsKernelLauncher( + get_ptr(input), get_ptr(permuted_input), + get_ptr(sorted_row_idx), get_ptr(inv_permuted_idx), + get_ptr(permuted_idx), get_ptr(expert_first_token_offset), + n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream); + }); +} + +void moe_permute( + const torch::stable::Tensor& input, // [n_token, hidden] + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] + const std::optional& expert_map, // [n_expert] + int64_t n_expert, int64_t n_local_expert, int64_t topk, + torch::stable::Tensor& permuted_input, // [permuted_size, hidden] + torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] + torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] + torch::stable::Tensor& permuted_idx) { // [permute_size] + moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, + n_local_expert, topk, permuted_input, + expert_first_token_offset, inv_permuted_idx, permuted_idx, + std::nullopt, std::nullopt, std::nullopt, std::nullopt); +} + +void moe_permute_with_scratch( + const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, int64_t n_expert, + int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& inv_permuted_idx, + torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& permuted_experts_id, + torch::stable::Tensor& sorted_row_idx, + torch::stable::Tensor& topk_ids_for_sort) { + moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, + n_local_expert, topk, permuted_input, + expert_first_token_offset, inv_permuted_idx, permuted_idx, + sort_workspace, permuted_experts_id, sorted_row_idx, + topk_ids_for_sort); +} + +void moe_unpermute( + const torch::stable::Tensor& + permuted_hidden_states, // [n_token * topk, hidden] + const torch::stable::Tensor& topk_weights, // [n_token, topk] + const torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] + const std::optional& + expert_first_token_offset, // [n_local_expert+1] + int64_t topk, + torch::stable::Tensor& hidden_states) { // [n_token, hidden] + STD_TORCH_CHECK( + permuted_hidden_states.scalar_type() == hidden_states.scalar_type(), + "permuted_hidden_states dtype must be same as hidden_states"); + + const torch::stable::accelerator::DeviceGuard device_guard( + hidden_states.get_device_index()); + auto n_token = hidden_states.size(0); + auto n_hidden = hidden_states.size(1); + auto stream = get_current_cuda_stream(hidden_states.get_device_index()); + + int64_t const* valid_ptr = nullptr; + if (expert_first_token_offset.has_value()) { + int n_local_expert = expert_first_token_offset.value().size(0) - 1; + valid_ptr = + get_ptr(expert_first_token_offset.value()) + n_local_expert; + } + + MOE_DISPATCH(hidden_states.scalar_type(), [&] { + finalizeMoeRoutingKernelLauncher( + get_ptr(permuted_hidden_states), + get_ptr(hidden_states), get_ptr(topk_weights), + get_ptr(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr, + stream); + }); +} + +template +__global__ void shuffleInputRowsKernel(const T* input, + const int32_t* dst2src_map, T* output, + int64_t num_src_rows, + int64_t num_dst_rows, int64_t num_cols) { + int64_t dest_row_idx = blockIdx.x; + int64_t const source_row_idx = dst2src_map[dest_row_idx]; + + if (blockIdx.x < num_dst_rows) { + // Load 128-bits per thread + constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8; + using DataElem = cutlass::Array; + + // Duplicate and permute rows + auto const* source_row_ptr = + reinterpret_cast(input + source_row_idx * num_cols); + auto* dest_row_ptr = + reinterpret_cast(output + dest_row_idx * num_cols); + + int64_t const start_offset = threadIdx.x; + int64_t const stride = blockDim.x; + int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD; + + for (int elem_index = start_offset; elem_index < num_elems_in_col; + elem_index += stride) { + dest_row_ptr[elem_index] = source_row_ptr[elem_index]; + } + } +} + +void shuffle_rows(const torch::stable::Tensor& input_tensor, + const torch::stable::Tensor& dst2src_map, + torch::stable::Tensor& output_tensor) { + STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(), + "Input and output tensors must have the same data type"); + + const torch::stable::accelerator::DeviceGuard device_guard( + output_tensor.get_device_index()); + auto stream = get_current_cuda_stream(output_tensor.get_device_index()); + const int64_t blocks = output_tensor.size(0); + const int64_t threads = 256; + const int64_t num_dest_rows = output_tensor.size(0); + const int64_t num_src_rows = input_tensor.size(0); + const int64_t num_cols = input_tensor.size(1); + + STD_TORCH_CHECK(!(num_cols % (128 / input_tensor.element_size() / 8)), + "num_cols must be divisible by 128 / " + "input_tensor.element_size() / 8"); + + MOE_DISPATCH(input_tensor.scalar_type(), [&] { + shuffleInputRowsKernel<<>>( + reinterpret_cast(input_tensor.const_data_ptr()), + reinterpret_cast(dst2src_map.const_data_ptr()), + reinterpret_cast(output_tensor.mutable_data_ptr()), + num_src_rows, num_dest_rows, num_cols); + }); +} + +#else + +int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, + int64_t n_expert) { + STD_TORCH_CHECK( + false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0"); +} + +void moe_permute(const torch::stable::Tensor& input, + const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, + int64_t n_expert, int64_t n_local_expert, int64_t topk, + torch::stable::Tensor& permuted_input, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& inv_permuted_idx, + torch::stable::Tensor& permuted_idx) { + STD_TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0"); +} + +void moe_permute_with_scratch( + const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, int64_t n_expert, + int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& inv_permuted_idx, + torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& permuted_experts_id, + torch::stable::Tensor& sorted_row_idx, + torch::stable::Tensor& topk_ids_for_sort) { + STD_TORCH_CHECK(false, + "moe_permute_with_scratch is not supported on CUDA < 12.0"); +} + +void moe_unpermute( + const torch::stable::Tensor& permuted_hidden_states, + const torch::stable::Tensor& topk_weights, + const torch::stable::Tensor& inv_permuted_idx, + const std::optional& expert_first_token_offset, + int64_t topk, torch::stable::Tensor& hidden_states) { + STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0"); +} + +#endif + +bool moe_permute_unpermute_supported() { +#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) + return true; +#else + return false; +#endif +} + +STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { + m.impl("moe_permute", TORCH_BOX(&moe_permute)); + m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch)); + m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute)); +} \ No newline at end of file diff --git a/csrc/moe/moe_wna16.cu b/csrc/libtorch_stable/moe/moe_wna16.cu similarity index 77% rename from csrc/moe/moe_wna16.cu rename to csrc/libtorch_stable/moe/moe_wna16.cu index 7b6a111c00a..9345a7c9f78 100644 --- a/csrc/moe/moe_wna16.cu +++ b/csrc/libtorch_stable/moe/moe_wna16.cu @@ -1,11 +1,14 @@ +#include -#include -#include -#include #include +#include +#include +#include +#include #include #include +#include "libtorch_stable/torch_utils.h" #include "moe_wna16_utils.h" #define DIVIDE(x, size) (((x) + (size) - 1) / (size)) @@ -263,7 +266,7 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output, } const int shared_mem_size = BLOCK_SIZE_M * BLOCK_SIZE_K * 2; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); kernel<<>>( input, output, b_qweight, b_scales, b_qzeros, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_pad, num_experts, @@ -271,17 +274,18 @@ void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output, BLOCK_SIZE_K, has_zp, mul_topk_weight); } -torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output, - torch::Tensor b_qweight, torch::Tensor b_scales, - std::optional b_qzeros, - std::optional topk_weights, - torch::Tensor sorted_token_ids, - torch::Tensor expert_ids, - torch::Tensor num_tokens_post_pad, int64_t top_k, - int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, - int64_t BLOCK_SIZE_K, int64_t bit) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - output.zero_(); +torch::stable::Tensor moe_wna16_gemm( + torch::stable::Tensor input, torch::stable::Tensor output, + torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales, + std::optional b_qzeros, + std::optional topk_weights, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, int64_t top_k, + int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K, + int64_t bit) { + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + torch::stable::zero_(output); const int num_experts = b_qweight.size(0); const int size_m = input.size(0); @@ -291,52 +295,56 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output, int64_t EM = sorted_token_ids.size(0); if (size_m <= BLOCK_SIZE_M) { - EM = min(EM, size_m * BLOCK_SIZE_M * top_k); + EM = std::min(EM, size_m * BLOCK_SIZE_M * top_k); } const int num_token_blocks = (EM + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; const uint32_t* b_qzeros_ptr; if (b_qzeros.has_value()) - b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr(); + b_qzeros_ptr = (const uint32_t*)b_qzeros.value().const_data_ptr(); const float* topk_weights_ptr = nullptr; if (topk_weights.has_value()) - topk_weights_ptr = (const float*)topk_weights.value().data_ptr(); + topk_weights_ptr = + (const float*)topk_weights.value().const_data_ptr(); int groups_per_block_row = BLOCK_SIZE_K / group_size; - TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8"); - TORCH_CHECK(size_k % BLOCK_SIZE_K == 0, - "size_k must divisible by BLOCK_SIZE_K"); - TORCH_CHECK(BLOCK_SIZE_K % group_size == 0, - "BLOCK_SIZE_K must divisible by group_size"); - TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64"); - TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 || - groups_per_block_row == 4 || groups_per_block_row == 8, - "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]"); + STD_TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8"); + STD_TORCH_CHECK(size_k % BLOCK_SIZE_K == 0, + "size_k must divisible by BLOCK_SIZE_K"); + STD_TORCH_CHECK(BLOCK_SIZE_K % group_size == 0, + "BLOCK_SIZE_K must divisible by group_size"); + STD_TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64"); + STD_TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 || + groups_per_block_row == 4 || groups_per_block_row == 8, + "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]"); - if (input.scalar_type() == at::ScalarType::Half) { + if (input.scalar_type() == torch::headeronly::ScalarType::Half) { run_moe_wna16_gemm( - (const half*)input.data_ptr(), - (half*)output.data_ptr(), - (const uint32_t*)b_qweight.data_ptr(), - (const half*)b_scales.data_ptr(), b_qzeros_ptr, - topk_weights_ptr, sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), num_tokens_post_pad.data_ptr(), - num_experts, group_size, num_token_blocks, top_k, size_m, size_n, - size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit, - b_qzeros.has_value(), topk_weights.has_value()); - } else if (input.scalar_type() == at::ScalarType::BFloat16) { + reinterpret_cast(input.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + (const uint32_t*)b_qweight.const_data_ptr(), + reinterpret_cast(b_scales.const_data_ptr()), b_qzeros_ptr, + topk_weights_ptr, sorted_token_ids.const_data_ptr(), + expert_ids.const_data_ptr(), + num_tokens_post_pad.const_data_ptr(), num_experts, group_size, + num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M, + BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(), + topk_weights.has_value()); + } else if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) { run_moe_wna16_gemm( - (const nv_bfloat16*)input.data_ptr(), - (nv_bfloat16*)output.data_ptr(), - (const uint32_t*)b_qweight.data_ptr(), - (const nv_bfloat16*)b_scales.data_ptr(), b_qzeros_ptr, - topk_weights_ptr, sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), num_tokens_post_pad.data_ptr(), - num_experts, group_size, num_token_blocks, top_k, size_m, size_n, - size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit, - b_qzeros.has_value(), topk_weights.has_value()); + reinterpret_cast(input.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + (const uint32_t*)b_qweight.const_data_ptr(), + reinterpret_cast(b_scales.const_data_ptr()), + b_qzeros_ptr, topk_weights_ptr, + sorted_token_ids.const_data_ptr(), + expert_ids.const_data_ptr(), + num_tokens_post_pad.const_data_ptr(), num_experts, group_size, + num_token_blocks, top_k, size_m, size_n, size_k, BLOCK_SIZE_M, + BLOCK_SIZE_N, BLOCK_SIZE_K, bit, b_qzeros.has_value(), + topk_weights.has_value()); } else { - TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16"); + STD_TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16"); } return output; } diff --git a/csrc/moe/moe_wna16_utils.h b/csrc/libtorch_stable/moe/moe_wna16_utils.h similarity index 100% rename from csrc/moe/moe_wna16_utils.h rename to csrc/libtorch_stable/moe/moe_wna16_utils.h diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu deleted file mode 100644 index fda9bc020da..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu +++ /dev/null @@ -1,69 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu - -#include -#include -#include "libtorch_stable/torch_utils.h" - -#include "cutlass_mxfp8_grouped_mm_launcher.cuh" - -void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& sfa, - const torch::stable::Tensor& sfb, - torch::stable::Tensor& d, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& blockscale_offsets) { -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); - STD_TORCH_CHECK(problem_sizes.size(1) == 3, - "problem_sizes must have 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"); - STD_TORCH_CHECK( - expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int, - "expert_offsets must be int32"); - STD_TORCH_CHECK( - blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int, - "blockscale_offsets must be int32"); - STD_TORCH_CHECK(a.dim() == 2, - "a must be a 2D tensor of shape (num_tokens, k)"); - STD_TORCH_CHECK(b.dim() == 3, - "b must be a 3D tensor of shape (num_experts, k, n)"); - STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0, - "k should align 128"); - STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128"); - STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major"); - STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major"); - - const torch::stable::accelerator::DeviceGuard device_guard( - a.get_device_index()); - auto stream = get_current_cuda_stream(a.get_device_index()); - if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) { - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< - cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, - blockscale_offsets, stream); - } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) { - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< - cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, - blockscale_offsets, stream); - } else { - STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); - } -#else - STD_TORCH_CHECK(false, - "No implemented cutlass_mxfp8_grouped_mm for " - "current device"); -#endif -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { - m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm)); -} diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh deleted file mode 100644 index 9fb1dbf8eef..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh +++ /dev/null @@ -1,141 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_functor.cuh - -#pragma once -#include - -#include "cute/tensor.hpp" -#include "cutlass/util/packed_stride.hpp" -#include "cutlass_mxfp8_grouped_mm_traits.cuh" - -namespace expert_specialization { - -using namespace cute; - -template -struct CutlassMxfp8GroupedMmOffsetFunctor { - using Gemm = typename GemmTraits::Gemm; - using ElementA = typename Gemm::ElementA; - using ElementB = typename Gemm::ElementB; - using ElementSF = typename GemmTraits::ElementSF; - using ElementD = typename GemmTraits::ElementOutput; - // Input - int* expert_offsets{nullptr}; - int* blockscale_offsets{nullptr}; - // Output - ElementA* a_base{nullptr}; - ElementB* b_base{nullptr}; - ElementSF* sfa_base{nullptr}; - ElementSF* sfb_base{nullptr}; - ElementD* d_base{nullptr}; - ElementA** a_offsets{nullptr}; - ElementB** b_offsets{nullptr}; - ElementSF** sfa_offsets{nullptr}; - ElementSF** sfb_offsets{nullptr}; - ElementD** d_offsets{nullptr}; - - CutlassMxfp8GroupedMmOffsetFunctor() = default; - CutlassMxfp8GroupedMmOffsetFunctor( - int* _expert_offsets, int* _blockscale_offsets, ElementA* _a_base, - ElementB* _b_base, ElementSF* _sfa_base, ElementSF* _sfb_base, - ElementD* _d_base, ElementA** _a_offsets, ElementB** _b_offsets, - ElementSF** _sfa_offsets, ElementSF** _sfb_offsets, ElementD** _d_offsets) - : expert_offsets{_expert_offsets}, - blockscale_offsets{_blockscale_offsets}, - a_base(_a_base), - b_base(_b_base), - sfa_base(_sfa_base), - sfb_base(_sfb_base), - d_base(_d_base), - a_offsets(_a_offsets), - b_offsets(_b_offsets), - sfa_offsets(_sfa_offsets), - sfb_offsets(_sfb_offsets), - d_offsets(_d_offsets) {} - - void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - int64_t expert_offset = static_cast(expert_offsets[expert_id]); - int64_t blockscale_offset = - static_cast(blockscale_offsets[expert_id]); - int64_t a_stride = expert_offset * k; - int64_t b_stride = expert_id * k * n; - int64_t d_stride = expert_offset * n; - int64_t sfa_stride = blockscale_offset * (k / 32); - int64_t sfb_stride = expert_id * n * (k / 32); - - a_offsets[expert_id] = a_base + a_stride; - b_offsets[expert_id] = b_base + b_stride; - sfa_offsets[expert_id] = sfa_base + sfa_stride; - sfb_offsets[expert_id] = sfb_base + sfb_stride; - d_offsets[expert_id] = d_base + d_stride; - } -}; - -template -struct CutlassMxfp8GroupedMmLayoutFunctor { - using Sm1xxBlkScaledConfig = typename GemmTraits::Sm1xxBlkScaledConfig; - using LayoutSFA = typename GemmTraits::LayoutSFA; - using LayoutSFB = typename GemmTraits::LayoutSFB; - LayoutSFA* layout_sfa_base{nullptr}; - LayoutSFB* layout_sfb_base{nullptr}; - - CutlassMxfp8GroupedMmLayoutFunctor() = default; - CutlassMxfp8GroupedMmLayoutFunctor(LayoutSFA* _layout_sfa_base, - LayoutSFB* _layout_sfb_base) - : layout_sfa_base(_layout_sfa_base), layout_sfb_base(_layout_sfb_base) {} - - void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - LayoutSFA* layout_sfa_ptr = layout_sfa_base + expert_id; - LayoutSFB* layout_sfb_ptr = layout_sfb_base + expert_id; - *layout_sfa_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA( - cute::make_shape(m, n, k, 1)); - *layout_sfb_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB( - cute::make_shape(m, n, k, 1)); - } -}; - -template -struct CutlassMxfp8GroupedMmStrideFunctor { - using StrideA = typename GemmTraits::StrideA; - using StrideB = typename GemmTraits::StrideB; - using StrideD = typename GemmTraits::StrideD; - StrideA* stride_A_base{nullptr}; - StrideB* stride_B_base{nullptr}; - StrideD* stride_D_base{nullptr}; - - CutlassMxfp8GroupedMmStrideFunctor() = default; - CutlassMxfp8GroupedMmStrideFunctor(StrideA* _stride_A_base, - StrideB* _stride_B_base, - StrideD* _stride_D_base) - : stride_A_base(_stride_A_base), - stride_B_base(_stride_B_base), - stride_D_base(_stride_D_base) {} - - void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - StrideA* stride_A = stride_A_base + expert_id; - StrideB* stride_B = stride_B_base + expert_id; - StrideD* stride_D = stride_D_base + expert_id; - *stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); - *stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); - *stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); - } -}; - -template -__global__ void cutlassMxfp8GroupedMmPreComputeKernel( - int* problem_sizes, OffsetFunctor offset_functor, - LayoutFunctor layout_functor, StrideFunctor stride_functor) { - int64_t expert_id = static_cast(threadIdx.x); - int m = problem_sizes[expert_id * 3 + 0]; - int n = problem_sizes[expert_id * 3 + 1]; - int k = problem_sizes[expert_id * 3 + 2]; - - offset_functor(expert_id, m, n, k); - layout_functor(expert_id, m, n, k); - stride_functor(expert_id, m, n, k); -} - -} // namespace expert_specialization \ No newline at end of file diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh deleted file mode 100644 index 82d6543b288..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh +++ /dev/null @@ -1,198 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh - -#pragma once - -#include -#include - -#include -#include -#include - -#include "cute/tensor.hpp" -#include "cutlass_mxfp8_grouped_mm_functor.cuh" -#include "cutlass_mxfp8_grouped_mm_traits.cuh" -#include "libtorch_stable/torch_utils.h" - -namespace expert_specialization { - -template -void cutlass_mxfp8_grouped_mm_pre_compute( - torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs, - torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs, - torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a, - torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d, - torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb, - const torch::stable::Tensor& a, const torch::stable::Tensor& b, - const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb, - const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) { - using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor; - using ElementA = typename OffsetFunctor::ElementA; - using ElementB = typename OffsetFunctor::ElementB; - using ElementSF = typename OffsetFunctor::ElementSF; - using ElementD = typename OffsetFunctor::ElementD; - - using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor; - using LayoutSFA = typename LayoutFunctor::LayoutSFA; - using LayoutSFB = typename LayoutFunctor::LayoutSFB; - - using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor; - using StrideA = typename StrideFunctor::StrideA; - using StrideB = typename StrideFunctor::StrideB; - using StrideD = typename StrideFunctor::StrideD; - - int num_experts = static_cast(expert_offsets.size(0)); - STD_TORCH_CHECK(num_experts <= 1024, - "Number of experts cannot exceed 1024, the maximum number of " - "threads per block."); - - OffsetFunctor offset_functor( - reinterpret_cast(expert_offsets.data_ptr()), - reinterpret_cast(blockscale_offsets.data_ptr()), - reinterpret_cast(a.data_ptr()), - reinterpret_cast(b.data_ptr()), - reinterpret_cast(sfa.data_ptr()), - reinterpret_cast(sfb.data_ptr()), - reinterpret_cast(d.data_ptr()), - reinterpret_cast(a_ptrs.data_ptr()), - reinterpret_cast(b_ptrs.data_ptr()), - reinterpret_cast(sfa_ptrs.data_ptr()), - reinterpret_cast(sfb_ptrs.data_ptr()), - reinterpret_cast(d_ptrs.data_ptr())); - LayoutFunctor layout_functor( - reinterpret_cast(layout_sfa.data_ptr()), - reinterpret_cast(layout_sfb.data_ptr())); - StrideFunctor stride_functor(reinterpret_cast(stride_a.data_ptr()), - reinterpret_cast(stride_b.data_ptr()), - reinterpret_cast(stride_d.data_ptr())); - cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>( - static_cast(problem_sizes.data_ptr()), offset_functor, - layout_functor, stride_functor); -} - -template -void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs, - const torch::stable::Tensor& b_ptrs, - const torch::stable::Tensor& sfa_ptrs, - const torch::stable::Tensor& sfb_ptrs, - const torch::stable::Tensor& d_ptrs, - const torch::stable::Tensor& stride_a, - const torch::stable::Tensor& stride_b, - const torch::stable::Tensor& stride_d, - const torch::stable::Tensor& layout_sfa, - const torch::stable::Tensor& layout_sfb, - const torch::stable::Tensor& problem_sizes, - cudaStream_t stream) { - using Gemm = typename GemmTraits::Gemm; - using ElementA = typename Gemm::ElementA; - using ElementB = typename Gemm::ElementB; - using ElementSF = typename GemmTraits::ElementSF; - using ElementD = typename GemmTraits::ElementOutput; - using StrideA = typename GemmTraits::StrideA; - using StrideB = typename GemmTraits::StrideB; - using StrideD = typename GemmTraits::StrideD; - using LayoutSFA = typename GemmTraits::LayoutSFA; - using LayoutSFB = typename GemmTraits::LayoutSFB; - using UnderlyingProblemShape = - typename GemmTraits::ProblemShape::UnderlyingProblemShape; - - cutlass::KernelHardwareInfo hw_info; - hw_info.device_id = d_ptrs.get_device_index(); - hw_info.sm_count = get_device_prop()->multiProcessorCount; - hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster; - hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster; - - int num_experts = static_cast(problem_sizes.size(0)); - - UnderlyingProblemShape* underlying_problem_shape = - reinterpret_cast(problem_sizes.data_ptr()); - - typename Gemm::Arguments arguments = { - cutlass::gemm::GemmUniversalMode::kGrouped, - {num_experts, underlying_problem_shape, nullptr}, - {reinterpret_cast(a_ptrs.data_ptr()), - reinterpret_cast(stride_a.data_ptr()), - reinterpret_cast(b_ptrs.data_ptr()), - reinterpret_cast(stride_b.data_ptr()), - reinterpret_cast(sfa_ptrs.data_ptr()), - reinterpret_cast(layout_sfa.data_ptr()), - reinterpret_cast(sfb_ptrs.data_ptr()), - reinterpret_cast(layout_sfb.data_ptr())}, - {{}, - nullptr, - nullptr, - reinterpret_cast(d_ptrs.data_ptr()), - reinterpret_cast(stride_d.data_ptr())}, - hw_info, - {} // Scheduler - }; - - Gemm gemm; - - auto can_implement_status = gemm.can_implement(arguments); - STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, - "Failed to implement GEMM"); - - size_t workspace_size = gemm.get_workspace_size(arguments); - torch::stable::Tensor workspace = torch::stable::empty( - {static_cast(workspace_size)}, - torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device()); - - auto status = gemm.initialize(arguments, workspace.data_ptr(), stream); - STD_TORCH_CHECK(status == cutlass::Status::kSuccess, - "Failed to initialize GEMM"); - - status = gemm.run(stream, nullptr, true); // Enable PDL - STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); -} - -template -void cutlass_mxfp8_grouped_mm_dispatch_out_dtype( - const torch::stable::Tensor& a, const torch::stable::Tensor& b, - const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb, - torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) { - int num_experts = static_cast(problem_sizes.size(0)); - auto device = a.device(); - - torch::stable::Tensor a_ptrs = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor b_ptrs = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor sfa_ptrs = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor sfb_ptrs = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor d_ptrs = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - - torch::stable::Tensor stride_a = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor stride_b = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor stride_d = torch::stable::empty( - num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::stable::Tensor layout_sfa = - torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int, - std::nullopt, device); - torch::stable::Tensor layout_sfb = - torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int, - std::nullopt, device); - - using GemmTraits = CutlassMxfp8GroupedMmGemmTraits; - cutlass_mxfp8_grouped_mm_pre_compute( - a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d, - layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets, - blockscale_offsets, stream); - cutlass_mxfp8_grouped_mm( - a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d, - layout_sfa, layout_sfb, problem_sizes, stream); -} - -} // namespace expert_specialization diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh deleted file mode 100644 index ed8cd7ce065..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh +++ /dev/null @@ -1,127 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_traits.cuh - -#pragma once - -// Misc -#include "cute/tensor.hpp" -#include "cutlass/arch/arch.h" -#include "cutlass/arch/mma.h" -#include "cutlass/cutlass.h" -#include "cutlass/detail/sm100_blockscaled_layout.hpp" -#include "cutlass/epilogue/dispatch_policy.hpp" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/layout/layout.h" -#include "cutlass/numeric_conversion.h" -#include "cutlass/numeric_size.h" - -// Collective Builder -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" -#include "cutlass/epilogue/thread/activation.h" -#include "cutlass/gemm/collective/collective_builder.hpp" - -// Integration -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" - -namespace expert_specialization { - -using namespace cute; - -// Different configs for 1SM and 2SM MMA kernel -struct MMA1SMConfig { - using MmaTileShape = Shape<_128, _128, _128>; - using KernelSchedule = - cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -const dim3 MMA1SMConfig::preferred_cluster(1, 4, 1); -const dim3 MMA1SMConfig::fallback_cluster(1, 2, 1); - -template -struct CutlassMxfp8GroupedMmGemmTraits { - using MMAConfig = _MMAConfig; - using ElementInput = cutlass::float_e4m3_t; - using ElementOutput = OutputDtype; - using ProblemShape = cutlass::gemm::GroupProblemShape>; - - // A matrix configuration - using ElementA = cutlass::mx_float8_t; - using LayoutA = cutlass::layout::RowMajor; - constexpr static int AlignmentA = 32; - - // B matrix configuration - using ElementB = cutlass::mx_float8_t; - using LayoutB = cutlass::layout::ColumnMajor; - constexpr static int AlignmentB = 32; - - // C/D matrix configuration - using ElementC = void; - using ElementD = ElementOutput; - using LayoutC = cutlass::layout::RowMajor; - using LayoutD = cutlass::layout::RowMajor; - constexpr static int AlignmentC = 128 / cutlass::sizeof_bits::value; - constexpr static int AlignmentD = 128 / cutlass::sizeof_bits::value; - using ElementAccumulator = float; - - static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest; - using CustomEVTIdentity = // acc - cutlass::epilogue::fusion::Sm90EVT< - cutlass::epilogue::fusion::Sm90Compute< - cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator, - RoundStyle>, - cutlass::epilogue::fusion::Sm90AccFetch>; - - // Core kernel configurations - using ArchTag = cutlass::arch::Sm100; - using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; - using StageCountType = cutlass::gemm::collective::StageCountAuto; - - // Runtime Cluster Shape - using ClusterShape = Shape; - - // Define Epilogue - using CollectiveEpilogue = - typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, OperatorClass, typename MMAConfig::MmaTileShape, - ClusterShape, Shape<_64, _64>, ElementAccumulator, ElementAccumulator, - ElementC, LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD, - typename MMAConfig::EpilogueSchedule, - CustomEVTIdentity>::CollectiveOp; - - // Define Mainloop - using CollectiveMainloop = - typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB, - LayoutB*, AlignmentB, ElementAccumulator, - typename MMAConfig::MmaTileShape, ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - typename MMAConfig::KernelSchedule>::CollectiveOp; - - // Define GemmKernel - using GemmKernel = - cutlass::gemm::kernel::GemmUniversal; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - - using ElementSF = typename Gemm::GemmKernel::ElementSF; - 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 Sm1xxBlkScaledConfig = - typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; -}; - -} // namespace expert_specialization \ No newline at end of file diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu deleted file mode 100644 index e075721c2a3..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu +++ /dev/null @@ -1,66 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu - -#include -#include -#include "libtorch_stable/torch_utils.h" - -#include "mxfp8_experts_quant.cuh" - -void mxfp8_experts_quant(const torch::stable::Tensor& input, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& blockscale_offsets, - torch::stable::Tensor& quant_output, - torch::stable::Tensor& scale_factor) { -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor"); - STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128"); - STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major"); - STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); - STD_TORCH_CHECK( - problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, - "problem_sizes must be int32"); - STD_TORCH_CHECK( - expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int, - "expert_offsets must be int32"); - STD_TORCH_CHECK( - blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int, - "blockscale_offsets must be int32"); - - auto groups = problem_sizes.size(0); - STD_TORCH_CHECK( - expert_offsets.dim() == 1 && expert_offsets.size(0) == groups, - "expert_offsets must be 1D and have size equal to the number of groups"); - STD_TORCH_CHECK( - blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups, - "blockscale_offsets must be 1D and have size equal to the number of " - "groups"); - - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) { - expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>( - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, - scale_factor); - } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) { - expert_specialization::launch_mxfp8_experts_quant<__half>( - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, - scale_factor); - } else { - STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); - } -#else - STD_TORCH_CHECK(false, - "No implemented mxfp8_experts_quant for " - "current device"); -#endif -} - -// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM -// is applied only under COMPILE_LANGUAGE:CUDA. -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { - m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant)); -} diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh deleted file mode 100644 index a57e00e76c3..00000000000 --- a/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh +++ /dev/null @@ -1,416 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh - -#pragma once -#include -#include -#include - -#include -#include -#include -#include - -#include - -#include "cute/tensor.hpp" -#include "libtorch_stable/torch_utils.h" - -namespace expert_specialization { - -using namespace cute; - -constexpr uint32_t THREAD_BLOCK_SIZE = 128; -constexpr uint32_t WARP_SIZE = 32; -constexpr int BLOCK_M = 128; -constexpr int BLOCK_K = 128; -using ThrLayout = Layout, Stride<_8, _1>>; -using ValLayout = Layout>; -using SfR2SThrLayout = Layout, Stride<_4, _1>>; -using SfR2SValLayout = Layout>; -using ScaleFactorTileLayout = - Layout, _4>, Stride, _1>>; - -// Fast reciprocal. -inline __device__ float reciprocal_approximate_ftz(float a) { - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -// Some code references TRT-LLM: -// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh -template -__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s, - FragmentD& fragment_d) { - using FragmentSLayout = typename FragmentS::layout_type; - using FragmentDLayout = typename FragmentD::layout_type; - FragmentSLayout fragment_s_layout; - FragmentDLayout fragment_d_layout; - static_assert(is_static::value && - size(fragment_s_layout) == 16); - static_assert(is_static::value && - size(fragment_d_layout) == 16); - - constexpr int eles_per_thr = 16; - using ValType = typename FragmentS::element_type; - using VecType = std::conditional_t, - __nv_bfloat162, __half2>; - VecType vec[8]; - // Assign vals - vec[0].x = fragment_s(Int<0>{}); - vec[0].y = fragment_s(Int<1>{}); - vec[1].x = fragment_s(Int<2>{}); - vec[1].y = fragment_s(Int<3>{}); - vec[2].x = fragment_s(Int<4>{}); - vec[2].y = fragment_s(Int<5>{}); - vec[3].x = fragment_s(Int<6>{}); - vec[3].y = fragment_s(Int<7>{}); - vec[4].x = fragment_s(Int<8>{}); - vec[4].y = fragment_s(Int<9>{}); - vec[5].x = fragment_s(Int<10>{}); - vec[5].y = fragment_s(Int<11>{}); - vec[6].x = fragment_s(Int<12>{}); - vec[6].y = fragment_s(Int<13>{}); - vec[7].x = fragment_s(Int<14>{}); - vec[7].y = fragment_s(Int<15>{}); - - auto local_max = __habs2(vec[0]); - for (int i = 1; i < eles_per_thr / 2; i++) { - local_max = __hmax2(__habs2(vec[i]), local_max); - } - local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max); - - // Get the final absolute maximum values. - float block_max(0.0f); - if constexpr (std::is_same_v) { - block_max = __bfloat162float(__hmax(local_max.x, local_max.y)); - } else { - block_max = __half2float(__hmax(local_max.x, local_max.y)); - } - // Get the SF (max value of the vector / max value of mxfp8). - float sf_val = block_max * reciprocal_approximate_ftz(448.0f); - // 8 bits representation of the SF. - uint8_t fp8_sf_val; - - __nv_fp8_e8m0 tmp_sf_val; - tmp_sf_val.__x = - __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf); - sf_val = static_cast(tmp_sf_val); - fp8_sf_val = tmp_sf_val.__x; - // Get the output scale (reciprocal of the SFValue). - float output_scale = - block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f; - - // Convert the input to float. - float2 fp2_vals[eles_per_thr / 2]; - -#pragma unroll - for (int i = 0; i < eles_per_thr / 2; i++) { - if constexpr (std::is_same_v) { - fp2_vals[i] = __half22float2(vec[i]); - } else { - fp2_vals[i] = __bfloat1622float2(vec[i]); - } - fp2_vals[i].x *= output_scale; - fp2_vals[i].y *= output_scale; - } - union { - uint8_t bytes[16]; - __nv_fp8x2_e4m3 elts[8]; - } u; - u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]); - u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]); - u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]); - u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]); - u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]); - u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]); - u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]); - u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]); - fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]); - fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]); - fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]); - fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]); - fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]); - fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]); - fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]); - fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]); - fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]); - fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]); - fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]); - fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]); - fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]); - fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]); - fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]); - fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]); - return fp8_sf_val; -} - -template -__inline__ __device__ void mxfp8_experts_quant_tile( - TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d, - TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m, - TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g, - TiledCopyR2S& tiled_copy_r2s) { - static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 && - size(get<1>(typename TensorS::layout_type{})) == 128 && - stride(get<1>(typename TensorS::layout_type{})) == 1); - static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 && - size(get<1>(typename TensorD::layout_type{})) == 128 && - stride(get<1>(typename TensorD::layout_type{})) == 1); - static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 && - size(get<1>(typename TensorP::layout_type{})) == 128); - static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 && - size(get<1>(typename TensorSharedSF::layout_type{})) == 4); - static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 && - size(get<1>(typename TensorSF::layout_type{})) == 4); - - using Tiler_MN = typename TiledCopyG2R::Tiler_MN; - auto tiler_mn = Tiler_MN{}; - static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128); - - auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn); - auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn); - auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn); - static_assert(size<2>(tiled_tensor_s) == 1); - static_assert(size<2>(tiled_tensor_p) == 1); - static_assert(size<2>(tiled_tensor_d) == 1); - auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s); - auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p); - auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d); - - using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN; - auto sf_tiler_mn = SF_Tiler_MN{}; - static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4); - - auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn); - auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn); - auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf); - auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf); - - constexpr int tile_loop_count = size<1>(tiled_tensor_s); - constexpr int rows_in_tile = 16; - // We don't need to clear shared memory - // clear(squeeze_tiled_tensor_shared_sf); -#pragma unroll 4 - for (int t = 0; t < tile_loop_count; t++) { - if (t * rows_in_tile >= m) { - break; - } - auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t)); - auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t)); - auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t)); - auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t)); - auto current_copy_tile_shared_sf = - tensor<0>(squeeze_tiled_tensor_shared_sf(_, t)); - - // Global to Register copy - auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x); - auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s); - auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p); - auto input_fragment = make_fragment_like(thr_tile_g2r_s); - - // Register to Global copy - auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x); - auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d); - auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p); - auto output_fragment = make_fragment_like(thr_tile_r2g_d); - - // Register to Shared copy - auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2); - auto thr_tile_r2s_shared_sf = - thr_copy_r2s.partition_D(current_copy_tile_shared_sf); - auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf); - - // CopyG2R & convert & CopyR2G - copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment); - uint8_t fp8_sf_val = - cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment); - copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d); - shared_sf_fragment[0] = fp8_sf_val; - - // Before first copy r2s, clear shared memory and wait previous group - if (t == 0 && threadIdx.x == 0) { - // Wait for the group to have completed reading from shared memory. - cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>()); - } - __syncthreads(); - - if (threadIdx.x % 2 == 0) { - copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf); - } - __syncthreads(); - } - - // Wait for shared memory writes to be visible to TMA engine. - cuda::ptx::fence_proxy_async(cuda::ptx::space_shared); // b) - __syncthreads(); - - if (threadIdx.x == 0) { - cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared, - squeeze_tiled_tensor_sf.data().get(), - squeeze_tiled_tensor_shared_sf.data().get(), 512); - // Wait for TMA transfer to have finished reading shared memory. - // Create a "bulk async-group" out of the previous bulk copy operation. - cuda::ptx::cp_async_bulk_commit_group(); - } - __syncthreads(); -} - -template -__global__ void mxfp8_experts_quant_kernel( - const T_IN* input, const int* problem_sizes, const int* expert_offsets, - const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output, - uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r, - TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - __shared__ __align__(512) uint8_t shared_memory[512]; - ScaleFactorTileLayout scale_factor_tile_layout{}; - auto scale_factor_shared = - make_tensor(make_smem_ptr(shared_memory), - scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1) - // TODO: Transform Groupwise Schedule into a more efficient Schedule - for (int g = 0; g < groups; g++) { - int m = problem_sizes[g * 3 + 0]; - int k = problem_sizes[g * 3 + 2]; - int64_t expert_offset = static_cast(expert_offsets[g]); - int64_t blockscale_offset = static_cast(blockscale_offsets[g]); - - auto input_tensor = make_tensor( - make_gmem_ptr(input + expert_offset * k), - make_layout(make_shape(m, k), - LayoutRight{})); // (M, K):(K, 1) half_t/bfloat16_t - - auto quant_output_tensor = make_tensor( - make_gmem_ptr(quant_output + expert_offset * k), - make_layout(make_shape(m, k), - LayoutRight{})); // (M, K):(K, 1) cutlass::float_e4m3_t - - auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32); - auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout, - scale_factor_shape, LayoutRight{}); - // layout<0>(layout<0>(scale_factor_layout)) (_32,_4):(_16,_4) -- static - // layout<1>(layout<0>(scale_factor_layout)) M_align_128 / 128 -- dynamic - // shape dynamic stride layout<0>(layout<1>(scale_factor_layout)) _4:_1 -- - // static layout<1>(layout<1>(scale_factor_layout)) (K / 32) / 4 : _512 -- - // dynamic shape static stride - - // Reshape to zipped layout for 1D indexing - auto zipped_scale_factor_layout = make_layout( - make_layout(layout<0>(layout<0>(scale_factor_layout)), - layout<0>(layout<1>(scale_factor_layout))), - make_layout( - layout<1>(layout<0>(scale_factor_layout)), - layout<1>(layout<1>( - scale_factor_layout)))); // (((_32,_4),_4),(M_align_128 / - // 128,(K / 32) / - // 4)):(((_16,_4),_1),(?,_512)) - - auto scale_factor_tensor = - make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)), - zipped_scale_factor_layout); - - // Used for cases where M is not divisible by 128 (most scenarios). - auto input_shape = shape(input_tensor); // (M, K):(K, 1) - auto identity_tensor = make_identity_tensor(input_shape); - auto predict_tensor = cute::lazy::transform( - identity_tensor, [&](auto c) { return elem_less(c, input_shape); }); - - // (_128, _128) - auto tiler = make_shape(Int{}, Int{}); - - auto tiled_input_tensor = zipped_divide( - input_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) - auto tiled_quant_output_tensor = - zipped_divide(quant_output_tensor, - tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) - auto tiled_predict_tensor = zipped_divide( - predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) - - auto total_tiles = - size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128) - decltype(total_tiles) blk_offset = blockIdx.x; - while (blk_offset < total_tiles) { - auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset)); - auto current_quant_output_tile = - tensor<0>(tiled_quant_output_tensor(_, blk_offset)); - auto current_predict_tile = - tensor<0>(tiled_predict_tensor(_, blk_offset)); - auto current_scale_factor_tile = - tensor<0>(scale_factor_tensor(_, blk_offset)); - - mxfp8_experts_quant_tile< - decltype(current_input_tile), decltype(current_predict_tile), - decltype(current_quant_output_tile), decltype(scale_factor_shared), - decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G, - TiledCopyR2S>(current_input_tile, current_predict_tile, - current_quant_output_tile, scale_factor_shared, - current_scale_factor_tile, m, tiled_copy_g2r, - tiled_copy_r2g, tiled_copy_r2s); - blk_offset += gridDim.x; - } - } -#endif -} - -template -void launch_mxfp8_experts_quant(const torch::stable::Tensor& input, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& blockscale_offsets, - torch::stable::Tensor& quant_output, - torch::stable::Tensor& scale_factor) { - ThrLayout thr_layout{}; - ValLayout val_layout{}; - SfR2SThrLayout r2s_thr_layout{}; - SfR2SValLayout r2s_val_layout{}; - - using CopyOpG2R = - UniversalCopy>; - using CopyAtomG2R = cute::Copy_Atom; - auto tiled_copy_g2r = cute::make_tiled_copy( - CopyAtomG2R{}, thr_layout, val_layout); // Tiler_MN: (16, 128) - - using CopyOpR2G = UniversalCopy< - cutlass::AlignedArray>; - using CopyAtomR2G = cute::Copy_Atom; - auto tiled_copy_r2g = cute::make_tiled_copy( - CopyAtomR2G{}, thr_layout, val_layout); // Tiler_MN: (16, 128) - - using CopyOpR2S = - UniversalCopy>; - using CopyAtomR2S = cute::Copy_Atom; - auto tiled_copy_r2s = cute::make_tiled_copy( - CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4) - - int max_active_blocks_per_sm = -1; - STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &max_active_blocks_per_sm, - mxfp8_experts_quant_kernel, - THREAD_BLOCK_SIZE, 0)); - - dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm, - 1, 1); - dim3 block(THREAD_BLOCK_SIZE, 1, 1); - int num_experts = static_cast(problem_sizes.size(0)); - auto stream = get_current_cuda_stream(input.get_device_index()); - mxfp8_experts_quant_kernel - <<>>( - reinterpret_cast(input.data_ptr()), - reinterpret_cast(problem_sizes.data_ptr()), - reinterpret_cast(expert_offsets.data_ptr()), - reinterpret_cast(blockscale_offsets.data_ptr()), - reinterpret_cast(quant_output.data_ptr()), - reinterpret_cast(scale_factor.data_ptr()), num_experts, - tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s); -} - -} // namespace expert_specialization \ No newline at end of file diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h b/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h new file mode 100644 index 00000000000..976233dd484 --- /dev/null +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h @@ -0,0 +1,60 @@ +#pragma once + +#include +#include +#include + +#define MOE_SWITCH(TYPE, ...) \ + const auto _st = (TYPE); \ + switch (_st) { \ + __VA_ARGS__ \ + default: \ + STD_TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \ + } + +#define MOE_DISPATCH_CASE(enum_type, ...) \ + case enum_type: { \ + using scalar_t = ScalarType2CudaType::type; \ + __VA_ARGS__(); \ + break; \ + } + +#define MOE_DISPATCH_FLOAT_CASE(...) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e5m2, __VA_ARGS__) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + MOE_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) + +#define MOE_DISPATCH(TYPE, ...) \ + MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__)) + +template +struct ScalarType2CudaType; + +template <> +struct ScalarType2CudaType { + using type = float; +}; +template <> +struct ScalarType2CudaType { + using type = half; +}; +template <> +struct ScalarType2CudaType { + using type = __nv_bfloat16; +}; +// uint8 for packed fp4 +template <> +struct ScalarType2CudaType { + using type = uint8_t; +}; +template <> +struct ScalarType2CudaType { + using type = __nv_fp8_e5m2; +}; +template <> +struct ScalarType2CudaType { + using type = __nv_fp8_e4m3; +}; \ No newline at end of file diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu similarity index 95% rename from csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu rename to csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index 2cc20032169..f5ec32c390f 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -1,5 +1,7 @@ +#include +#include -#include "moe_permute_unpermute_kernel.h" +#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" // moe_permute kernels require at least CUDA 12.0 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) @@ -48,9 +50,10 @@ void CubKeyValueSorter::run(void* workspace, size_t const workspace_size, size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs, num_experts_); size_t actual_ws_size = workspace_size; - TORCH_CHECK(expected_ws_size <= workspace_size, - "[CubKeyValueSorter::run] The allocated workspace is too small " - "to run this problem."); + STD_TORCH_CHECK( + expected_ws_size <= workspace_size, + "[CubKeyValueSorter::run] The allocated workspace is too small " + "to run this problem."); cub::DeviceRadixSort::SortPairs(workspace, actual_ws_size, keys_in, keys_out, values_in, values_out, num_key_value_pairs, 0, num_bits_, stream); diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h similarity index 89% rename from csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h rename to csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index fe44d301559..89c278a4ed4 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -2,23 +2,24 @@ // reference from tensorrt_llm moe kernel implementation archive in // https://github.com/BBuf/tensorrt-llm-moe/tree/master -#include -#include -#include "dispatch.h" +#include + #include #include #include -#include "cutlass/numeric_size.h" + #include "cutlass/array.h" +#include "cutlass/numeric_size.h" +#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h" template -inline T* get_ptr(torch::Tensor& t) { - return reinterpret_cast(t.data_ptr()); +inline T* get_ptr(torch::stable::Tensor& t) { + return reinterpret_cast(t.mutable_data_ptr()); } template -inline const T* get_ptr(const torch::Tensor& t) { - return reinterpret_cast(t.data_ptr()); +inline const T* get_ptr(const torch::stable::Tensor& t) { + return reinterpret_cast(t.const_data_ptr()); } class CubKeyValueSorter { diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl similarity index 100% rename from csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl rename to csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu similarity index 83% rename from csrc/moe/topk_softmax_kernels.cu rename to csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 57461a044f9..b4bcd9479e9 100644 --- a/csrc/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -17,11 +17,16 @@ * limitations under the License. */ #include -#include -#include -#include -#include "../cuda_compat.h" + +#include +#include +#include +#include +#include + +#include "../../cuda_compat.h" #include "../cub_helpers.h" +#include "libtorch_stable/torch_utils.h" #ifndef USE_ROCM #include @@ -168,7 +173,8 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int start_expert, const int end_expert, const bool renormalize, - const float* bias) + const float* bias, + const double routed_scaling_factor) { using cub_kvp = cub::KeyValuePair; @@ -236,14 +242,16 @@ __launch_bounds__(TPB) __global__ void moeTopK( __syncthreads(); } - // Renormalize the k weights for this row to sum to 1, if requested. - if (renormalize) { - if (threadIdx.x == 0) { + // Apply renormalization and routed scaling factor to final weights. + if (threadIdx.x == 0) { + float scale = static_cast(routed_scaling_factor); + if (renormalize) { const float denom = selected_sum > 0.f ? selected_sum : 1.f; - for (int k_idx = 0; k_idx < k; ++k_idx) { - const int idx = k * block_row + k_idx; - output[idx] = output[idx] / denom; - } + scale /= denom; + } + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] = output[idx] * scale; } } } @@ -269,7 +277,7 @@ template || std::is_same_v || std::is_same_v, @@ -565,17 +573,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ } } - // Renormalize the k weights for this row to sum to 1, if requested. - if (renormalize) { - if (thread_group_idx == 0) - { - const float denom = selected_sum > 0.f ? selected_sum : 1.f; - for (int k_idx = 0; k_idx < k; ++k_idx) - { - const int idx = k * thread_row + k_idx; - output[idx] = output[idx] / denom; - } - } + // 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; + } } } @@ -597,7 +605,7 @@ struct TopkConstants template void topkGatingLauncherHelper(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, - const float* bias, cudaStream_t stream) + const float* bias, const double routed_scaling_factor, cudaStream_t stream) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = detail::TopkConstants; @@ -608,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); topkGating<<>>( - input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias); + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor); } #ifndef USE_ROCM @@ -619,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); + bias, routed_scaling_factor, stream); #else #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -627,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); \ + bias, routed_scaling_factor, stream); \ } else if (WARP_SIZE == 32) { \ topkGatingLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); \ + bias, routed_scaling_factor, stream); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -652,6 +660,7 @@ void topkGatingKernelLauncher( const int topk, const bool renormalize, const float* bias, + const double routed_scaling_factor, cudaStream_t stream) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; @@ -713,7 +722,7 @@ void topkGatingKernelLauncher( break; #endif default: { - TORCH_CHECK(workspace != nullptr, + STD_TORCH_CHECK(workspace != nullptr, "workspace must be provided for num_experts that are not a power of 2 or multiple of 64."); static constexpr int TPB = 256; if constexpr (SF == SCORING_SOFTMAX) { @@ -723,11 +732,11 @@ void topkGatingKernelLauncher( moeSigmoid<<>>( gating_output, nullptr, workspace, num_experts); } else { - TORCH_CHECK(false, "Unsupported scoring func"); + STD_TORCH_CHECK(false, "Unsupported scoring func"); } moeTopK<<>>( workspace, nullptr, topk_weights, topk_indices, token_expert_indices, - num_experts, topk, 0, num_experts, renormalize, bias); + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor); } } } @@ -738,63 +747,66 @@ void topkGatingKernelLauncher( template void dispatch_topk_launch( - torch::Tensor& gating_output, - torch::Tensor& topk_weights, - torch::Tensor& topk_indices, - torch::Tensor& token_expert_indices, - torch::Tensor& softmax_workspace, + torch::stable::Tensor& gating_output, + torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& softmax_workspace, int num_tokens, int num_experts, int topk, bool renormalize, - std::optional bias, + std::optional bias, + double routed_scaling_factor, cudaStream_t stream) { const float* bias_ptr = nullptr; if (bias.has_value()) { - const torch::Tensor& bias_tensor = bias.value(); - TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "bias tensor must be float32"); - TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D"); - TORCH_CHECK(bias_tensor.size(0) == num_experts, "bias size mismatch, expected: ", num_experts); - TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous"); - bias_ptr = bias_tensor.data_ptr(); + const torch::stable::Tensor& bias_tensor = bias.value(); + STD_TORCH_CHECK(bias_tensor.scalar_type() == torch::headeronly::ScalarType::Float, + "bias tensor must be float32"); + STD_TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D"); + STD_TORCH_CHECK(bias_tensor.size(0) == num_experts, + "bias size mismatch, expected: ", num_experts); + STD_TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous"); + bias_ptr = bias_tensor.const_data_ptr(); } - if (topk_indices.scalar_type() == at::ScalarType::Int) { + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { vllm::moe::topkGatingKernelLauncher( - reinterpret_cast(gating_output.data_ptr()), - topk_weights.data_ptr(), - topk_indices.data_ptr(), - token_expert_indices.data_ptr(), - softmax_workspace.data_ptr(), + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); - } else if (topk_indices.scalar_type() == at::ScalarType::UInt32) { + bias_ptr, routed_scaling_factor, stream); + } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { vllm::moe::topkGatingKernelLauncher( - reinterpret_cast(gating_output.data_ptr()), - topk_weights.data_ptr(), - topk_indices.data_ptr(), - token_expert_indices.data_ptr(), - softmax_workspace.data_ptr(), + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); + bias_ptr, routed_scaling_factor, stream); } else { - TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long); + STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); vllm::moe::topkGatingKernelLauncher( - reinterpret_cast(gating_output.data_ptr()), - topk_weights.data_ptr(), - topk_indices.data_ptr(), - token_expert_indices.data_ptr(), - softmax_workspace.data_ptr(), + reinterpret_cast(gating_output.const_data_ptr()), + topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_data_ptr(), + softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); + bias_ptr, routed_scaling_factor, stream); } } void topk_softmax( - 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] + torch::stable::Tensor& topk_weights, // [num_tokens, topk] + torch::stable::Tensor& topk_indices, // [num_tokens, topk] + torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] + torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, - std::optional bias) + std::optional bias) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -804,35 +816,37 @@ void topk_softmax( const bool needs_workspace = !is_pow_2 || num_experts > 256; const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; - const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float); - torch::Tensor softmax_workspace = torch::empty({workspace_size}, workspace_options); + torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(gating_output.get_device_index()); + auto softmax_workspace = torch::stable::new_empty( + gating_output, {workspace_size}, torch::headeronly::ScalarType::Float); - if (gating_output.scalar_type() == at::ScalarType::Float) { + if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); - } else if (gating_output.scalar_type() == at::ScalarType::Half) { + bias, 1.0, stream); + } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); - } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) { + bias, 1.0, stream); + } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, 1.0, stream); } else { - TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); + STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } } void topk_sigmoid( - 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] + torch::stable::Tensor& topk_weights, // [num_tokens, topk] + torch::stable::Tensor& topk_indices, // [num_tokens, topk] + torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] + torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, - std::optional bias) + std::optional bias, + double routed_scaling_factor) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -842,24 +856,25 @@ void topk_sigmoid( const bool needs_workspace = !is_pow_2 || num_experts > 256; const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; - const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - const auto workspace_options = gating_output.options().dtype(at::ScalarType::Float); - torch::Tensor workspace = torch::empty({workspace_size}, workspace_options); + torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(gating_output.get_device_index()); + auto workspace = torch::stable::new_empty( + gating_output, {workspace_size}, torch::headeronly::ScalarType::Float); - if (gating_output.scalar_type() == at::ScalarType::Float) { + if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); - } else if (gating_output.scalar_type() == at::ScalarType::Half) { + bias, routed_scaling_factor, stream); + } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); - } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) { + bias, routed_scaling_factor, stream); + } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, routed_scaling_factor, stream); } else { - TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); + STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } } diff --git a/csrc/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu similarity index 87% rename from csrc/moe/topk_softplus_sqrt_kernels.cu rename to csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index d5bb8edadc6..095a7667831 100644 --- a/csrc/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -18,11 +18,16 @@ * limitations under the License. */ #include -#include -#include -#include -#include "../cuda_compat.h" + +#include +#include +#include +#include +#include + +#include "../../cuda_compat.h" #include "../cub_helpers.h" +#include "libtorch_stable/torch_utils.h" #ifndef USE_ROCM #include #include @@ -168,7 +173,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ float row_chunk[VPT]; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert @@ -295,7 +300,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif return; } else { @@ -420,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } } @@ -618,7 +623,7 @@ void topkGatingSoftplusSqrtKernelLauncher( LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW); break; default: { - TORCH_CHECK(false, "Unsupported expert number: ", num_experts); + STD_TORCH_CHECK(false, "Unsupported expert number: ", num_experts); } } } @@ -628,100 +633,109 @@ void topkGatingSoftplusSqrtKernelLauncher( 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 ComputeType* gating_output, torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, int num_tokens, + int num_experts, int topk, bool renormalize, double routed_scaling_factor, + const std::optional& correction_bias, + const std::optional& input_ids, + const std::optional& tid2eid, cudaStream_t stream) { const float* bias_ptr = nullptr; if (correction_bias.has_value()) { - bias_ptr = correction_bias.value().data_ptr(); + bias_ptr = correction_bias.value().const_data_ptr(); } bool use_hash = false; if (tid2eid.has_value()) { - TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE"); + STD_TORCH_CHECK(input_ids.has_value(), + "input_ids is required for hash MoE"); use_hash = true; } - if (topk_indices.scalar_type() == at::ScalarType::Int) { + if (topk_indices.scalar_type() == torch::headeronly::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(); + input_ids_ptr = input_ids.value().const_data_ptr(); + tid2eid_ptr = tid2eid.value().const_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) { + gating_output, topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_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() == + torch::headeronly::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(); + input_ids_ptr = input_ids.value().const_data_ptr(); + tid2eid_ptr = tid2eid.value().const_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); + gating_output, topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_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); + STD_TORCH_CHECK(topk_indices.scalar_type() == + torch::headeronly::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(); + input_ids_ptr = input_ids.value().const_data_ptr(); + tid2eid_ptr = tid2eid.value().const_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); + gating_output, topk_weights.mutable_data_ptr(), + topk_indices.mutable_data_ptr(), + token_expert_indices.mutable_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] + torch::stable::Tensor& topk_weights, // [num_tokens, topk] + torch::stable::Tensor& topk_indices, // [num_tokens, topk] + torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] + torch::stable::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 std::optional& correction_bias, + const std::optional& input_ids, + const std::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(); + const torch::stable::accelerator::DeviceGuard guard( + gating_output.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(gating_output.get_device_index()); - if (gating_output.scalar_type() == at::ScalarType::Float) { + if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_softplus_sqrt_launch( - gating_output.data_ptr(), topk_weights, topk_indices, + gating_output.const_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) { + } else if (gating_output.scalar_type() == + torch::headeronly::ScalarType::Half) { dispatch_topk_softplus_sqrt_launch<__half>( - reinterpret_cast(gating_output.data_ptr()), + reinterpret_cast(gating_output.const_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) { + } else if (gating_output.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( - reinterpret_cast( - gating_output.data_ptr()), + reinterpret_cast(gating_output.const_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()); + STD_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/libtorch_stable/moe/torch_bindings.cpp similarity index 81% rename from csrc/moe/torch_bindings.cpp rename to csrc/libtorch_stable/moe/torch_bindings.cpp index 99230f03b4b..ba5b9b896f1 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -1,32 +1,30 @@ #include "core/registration.h" #include "moe_ops.h" -TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { +#include + +STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { // Apply topk softmax to the gating outputs. m.def( "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? " "bias) -> ()"); - m.impl("topk_softmax", torch::kCUDA, &topk_softmax); // Apply topk sigmoid to the gating outputs. m.def( "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! " - "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? " - "bias) -> ()"); - m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid); + "token_expert_indices, Tensor gating_output, bool renormalize, " + "Tensor? bias, float routed_scaling_factor) -> ()"); 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); // Calculate the result of moe by summing up the partial results // from all selected experts. m.def("moe_sum(Tensor input, Tensor! output) -> ()"); - m.impl("moe_sum", torch::kCUDA, &moe_sum); // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size. @@ -36,7 +34,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { " Tensor! experts_ids," " Tensor! num_tokens_post_pad," " Tensor? maybe_expert_map) -> ()"); - m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size); // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size, but for the batched case. @@ -46,8 +43,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { " Tensor! sorted_token_ids," " Tensor! experts_ids," " Tensor! num_tokens_post_pad) -> ()"); - m.impl("batched_moe_align_block_size", torch::kCUDA, - &batched_moe_align_block_size); // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size. @@ -64,8 +59,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { " Tensor !adapter_enabled," " Tensor !lora_ids," " Tensor? maybe_expert_map) -> () "); - m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size); - #ifndef USE_ROCM m.def( "moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, " @@ -75,8 +68,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { "int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, " "int bit) -> Tensor"); - m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm); - m.def( "moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none," "Tensor! b_q_weight, Tensor? b_bias_or_none," @@ -118,14 +109,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { m.def( "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> " "int"); - m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported); - m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size); // Row shuffle for MoE m.def( "shuffle_rows(Tensor input_tensor, Tensor dst2src_map, Tensor! " "output_tensor) -> ()"); - m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows); // Apply grouped topk routing to select experts. m.def( @@ -133,7 +121,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { "topk_group, int topk, bool renormalize, float " "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, " "Tensor)"); - m.impl("grouped_topk", torch::kCUDA, &grouped_topk); // DeepSeek V3 optimized router GEMM for SM90+ m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); @@ -141,4 +128,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { #endif } -REGISTER_EXTENSION(TORCH_EXTENSION_NAME) +STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { + m.impl("topk_softmax", TORCH_BOX(&topk_softmax)); + m.impl("topk_sigmoid", TORCH_BOX(&topk_sigmoid)); + m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt)); + m.impl("moe_sum", TORCH_BOX(&moe_sum)); + m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size)); + m.impl("batched_moe_align_block_size", + TORCH_BOX(&batched_moe_align_block_size)); + m.impl("moe_lora_align_block_size", TORCH_BOX(&moe_lora_align_block_size)); +#ifndef USE_ROCM + m.impl("moe_wna16_gemm", TORCH_BOX(&moe_wna16_gemm)); + m.impl("shuffle_rows", TORCH_BOX(&shuffle_rows)); + m.impl("grouped_topk", TORCH_BOX(&grouped_topk)); +#endif +} + +#ifndef USE_ROCM +// Primitive-only ops have no tensor to dispatch on. +STABLE_TORCH_LIBRARY_IMPL(_moe_C, CompositeExplicitAutograd, m) { + m.impl("moe_permute_unpermute_supported", + TORCH_BOX(&moe_permute_unpermute_supported)); + m.impl("moe_permute_sort_workspace_size", + TORCH_BOX(&moe_permute_sort_workspace_size)); +} +#endif + +REGISTER_EXTENSION(_moe_C_stable_libtorch) diff --git a/csrc/libtorch_stable/ngram_embedding_kernels.cu b/csrc/libtorch_stable/ngram_embedding_kernels.cu new file mode 100644 index 00000000000..a4da62f9acd --- /dev/null +++ b/csrc/libtorch_stable/ngram_embedding_kernels.cu @@ -0,0 +1,96 @@ +// N-gram embedding index kernel for LongCat-Flash (n-gram embedding variant). +// +// Adapted from SGLang: +// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/ngram_embedding.cuh +// +// For each position, computes the hashed n-gram embedding ids that index the +// concatenated embedder table. Integer tensors are int32 except ``row_indices`` +// (int64); the token table is ``[max_running_reqs, max_context_len]`` int32, +// where a negative entry marks an ignored token (e.g. an EOS boundary). + +#include "torch_utils.h" + +#include "ops.h" + +#include + +namespace vllm::ngram_embedding { + +constexpr int kBlockThreads = 256; + +__global__ void ComputeNGramIdsKernel( + int batch_size, int ne_n, int ne_k, + int* ne_weights, // [ne_n-1, ne_k, ne_n] + int* ne_mods, // [ne_n-1, ne_k] + int* exclusive_ne_embedder_size_sums, // [(ne_n-1)*ne_k + 1] + int* exclusive_req_len_sums, // [batch_size + 1] + int* ne_token_table, // [max_running_reqs, max_context_len] + int max_context_len, + const int64_t* __restrict__ row_indices, // [batch_size] + int* column_starts, // [batch_size] + int* n_gram_ids // [token_num, (ne_n-1)*ne_k] +) { + const int req_id = blockIdx.x % batch_size; + const int config_id = (blockIdx.x - req_id) / batch_size; + // n and k are offset from their physical meaning: n = real_n - 2, k = real_k + // - 1 (they index into ne_weights / ne_mods). + const int k = config_id % ne_k; + const int n = (config_id - config_id % ne_k) / ne_k; + const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n; + const int ne_mod = ne_mods[n * ne_k + k]; + for (int i = exclusive_req_len_sums[req_id] + threadIdx.x; + i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) { + uint64_t n_gram_id = 0; + const int64_t current_token_offset = i - exclusive_req_len_sums[req_id]; + const int64_t req_token_table_index = + row_indices[req_id] * static_cast(max_context_len); + const int64_t current_token_table_index = + req_token_table_index + column_starts[req_id] + current_token_offset; + for (int j = 0; j < n + 2; j++) { + if (current_token_table_index - j < req_token_table_index) { + break; // outside this request's range + } + if (ne_token_table[current_token_table_index - j] < 0) { + break; // ignored token + } + const uint64_t term = + (uint64_t)ne_token_table[current_token_table_index - j] * + (uint64_t)ne_weights[ne_weight_base_idx + j]; + n_gram_id += term % ne_mod; + } + n_gram_id %= ne_mod; + n_gram_id += exclusive_ne_embedder_size_sums[n * ne_k + k]; + n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id); + } +} + +} // namespace vllm::ngram_embedding + +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids) { + const int batch_size = static_cast(exclusive_req_len_sums.size(0) - 1); + const int max_context_len = static_cast(ne_token_table.size(1)); + const int num_configs = (static_cast(ne_n) - 1) * static_cast(ne_k); + const int grid_size = num_configs * batch_size; + if (grid_size <= 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + ne_weights.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + vllm::ngram_embedding::ComputeNGramIdsKernel<<< + grid_size, vllm::ngram_embedding::kBlockThreads, 0, stream>>>( + batch_size, static_cast(ne_n), static_cast(ne_k), + ne_weights.mutable_data_ptr(), + ne_mods.mutable_data_ptr(), + exclusive_ne_embedder_size_sums.mutable_data_ptr(), + exclusive_req_len_sums.mutable_data_ptr(), + ne_token_table.mutable_data_ptr(), max_context_len, + row_indices.const_data_ptr(), + column_starts.mutable_data_ptr(), + n_gram_ids.mutable_data_ptr()); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 6ebec954497..ae274226fe2 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -2,6 +2,25 @@ #include #include +#include + +#include +#include +#include + +#include + +inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { + // Ensure tensor is on CUDA + STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); + + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); + + /// Create a new tensor from the raw data pointer + return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), + tensor.device(), tensor.scalar_type()); +} void per_token_group_quant_fp8(const torch::stable::Tensor& input, torch::stable::Tensor& output_q, @@ -162,13 +181,12 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, // AllSpark ops: declarations are in the source files // (allspark_repack.cu and allspark_qgemm_w8a16.cu) -// TODO: Move this out once ROCm upgrade their torch to 2.11. -// CPU tensor -> CUDA UVA view (shared CUDA) +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) torch::stable::Tensor get_cuda_view_from_cpu_tensor( torch::stable::Tensor& cpu_tensor); -#endif - // Attention kernels (shared CUDA/ROCm) void merge_attn_states( torch::stable::Tensor& output, @@ -185,11 +203,12 @@ torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, // Layernorm kernels (shared CUDA/ROCm) void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - torch::stable::Tensor& weight, double epsilon); + std::optional weight, double epsilon); void fused_add_rms_norm(torch::stable::Tensor& input, torch::stable::Tensor& residual, - torch::stable::Tensor& weight, double epsilon); + std::optional weight, + double epsilon); // Layernorm-quant kernels (shared CUDA/ROCm) void rms_norm_static_fp8_quant(torch::stable::Tensor& out, @@ -268,10 +287,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( int64_t cache_block_size); #ifndef USE_ROCM -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps); std::tuple minimax_allreduce_rms_qk(torch::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, @@ -281,6 +296,25 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv, int64_t const nranks, double const eps); #endif +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs +// the index branch and scatters k/v/index_k into their paged caches. +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, int64_t num_heads, + int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional index_q_out, + const std::string& kv_cache_dtype, bool skip_index_branch); + // Sampler kernels (shared CUDA/ROCm) void apply_repetition_penalties_( torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, @@ -304,6 +338,14 @@ void persistent_topk(const torch::stable::Tensor& logits, torch::stable::Tensor& workspace, int64_t k, int64_t max_seq_len); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); +#endif + void selective_scan_fwd( const torch::stable::Tensor& u, const torch::stable::Tensor& delta, const torch::stable::Tensor& A, const torch::stable::Tensor& B, @@ -346,7 +388,20 @@ void free_shared_buffer(int64_t buffer); // Activation kernels (shared CUDA/ROCm) void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void silu_and_mul_clamp(torch::stable::Tensor& out, - torch::stable::Tensor& input, double limit); + torch::stable::Tensor& input, double limit, + double alpha = 1.0, double beta = 0.0); + +void silu_and_mul_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& scale); + +void persistent_masked_m_silu_mul_quant( + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] + bool use_ue8m0); + void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_tanh_and_mul(torch::stable::Tensor& out, @@ -359,6 +414,8 @@ void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); +void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); + // INT8 quantization kernels (shared CUDA/ROCm) void static_scaled_int8_quant(torch::stable::Tensor& out, torch::stable::Tensor const& input, @@ -397,61 +454,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, int64_t bit); -// GGML kernels (shared CUDA/ROCm) -torch::stable::Tensor ggml_dequantize( - torch::stable::Tensor W, int64_t type, int64_t m, int64_t n, - std::optional const& dtype); - -torch::stable::Tensor ggml_mul_mat_vec_a8(torch::stable::Tensor W, - torch::stable::Tensor X, int64_t type, - int64_t row); - -torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W, - torch::stable::Tensor X, int64_t type, - int64_t row); - -torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X, - torch::stable::Tensor W, - torch::stable::Tensor sorted_token_ids, - torch::stable::Tensor expert_ids, - torch::stable::Tensor num_tokens_post_padded, - int64_t type, int64_t row, int64_t top_k, - int64_t tokens); - -torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X, - torch::stable::Tensor W, - torch::stable::Tensor topk_ids, - int64_t top_k, int64_t type, int64_t row, - int64_t tokens); - -int64_t ggml_moe_get_block_size(int64_t type); - -void paged_attention_v1( - torch::stable::Tensor& out, torch::stable::Tensor& query, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, - torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - -void paged_attention_v2( - torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, - torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, - torch::stable::Tensor& query, torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - // Cache ops (shared CUDA/ROCm) void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, int64_t block_size_in_bytes, @@ -554,3 +556,12 @@ void cp_gather_indexer_k_quant_cache( // quant_block_size * 4] const torch::stable::Tensor& block_table, // [batch_size, num_blocks] const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 6b25dc9940e..85618feeb8a 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -11,6 +11,8 @@ #include #include +#include "topk_histogram_4096.cuh" + namespace vllm { namespace persistent { @@ -935,8 +937,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) } // namespace persistent // ============================================================================ -// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only. -// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent). +// ============================================================================ +// Optimized FilteredTopK — single CTA per row for bs > 32. +// Kept with persistent_topk so the portable fallback owns the non-cluster path. +// ============================================================================ +namespace filtered_topk { + +namespace hist4096 = topk_histogram_4096; + +// ============================================================================ +// FilteredTopK — single CTA per row for bs > 32 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ @@ -963,13 +973,6 @@ struct vec_t { 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 @@ -1013,7 +1016,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = * \tparam IdType Index type (int32_t) * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) */ -template +template __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) FilteredTopKUnifiedKernel(const DType* __restrict__ input, IdType* __restrict__ output, @@ -1042,6 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) return; } + // Short path + if (length <= 32768) { + extern __shared__ uint8_t _smem_reg[]; + if constexpr (UsePredicatedShortLoads) { + hist4096::histogram_4096_topk_predicated(score, dst, length, + _smem_reg); + } else { + hist4096::histogram_4096_topk(score, dst, length, + _smem_reg); + } + return; + } + // Static shared memory alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; alignas(128) __shared__ int s_counter; @@ -1285,14 +1302,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, 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; \ +#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) @@ -1306,6 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, return cudaSuccess; } +} // namespace filtered_topk + +template +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, + cudaStream_t stream = 0) { + return filtered_topk::FilteredTopKRaggedTransform( + input, output_indices, lengths, num_rows, top_k_val, max_len, stream); +} + } // namespace vllm #endif // PERSISTENT_TOPK_CUH_ diff --git a/csrc/quantization/activation_kernels.cu b/csrc/libtorch_stable/quantization/activation_kernels.cu similarity index 87% rename from csrc/quantization/activation_kernels.cu rename to csrc/libtorch_stable/quantization/activation_kernels.cu index 8cc645c33e2..822a41969e7 100644 --- a/csrc/quantization/activation_kernels.cu +++ b/csrc/libtorch_stable/quantization/activation_kernels.cu @@ -1,16 +1,12 @@ -#include -#include -#include +#include "libtorch_stable/torch_utils.h" #include -#include "core/math.hpp" -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#include "libtorch_stable/core/math.hpp" +#include "cuda_compat.h" +#include "libtorch_stable/dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" -#include - #ifndef USE_ROCM #include #include @@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3; #endif #endif -#include "core/registration.h" namespace vllm { template @@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel( } // namespace vllm // Launch activation, gating, and quantize kernel. -#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \ - int d = input.size(-1) / 2; \ - int64_t num_tokens = input.numel() / input.size(-1); \ - dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \ - dim3 block(std::min(d, 512)); \ - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ - VLLM_DISPATCH_FLOATING_TYPES( \ - input.scalar_type(), "act_and_mul_kernel", [&] { \ - VLLM_DISPATCH_FP8_TYPES( \ - out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \ - vllm::act_and_mul_quant_kernel, \ - fp8_t> \ - <<>>(out.data_ptr(), \ - input.data_ptr(), \ - scale.data_ptr(), d); \ - }); \ +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \ + dim3 block(std::min(d, 512)); \ + 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_FLOATING_TYPES( \ + input.scalar_type(), "act_and_mul_kernel", [&] { \ + VLLM_STABLE_DISPATCH_FP8_TYPES( \ + out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \ + vllm::act_and_mul_quant_kernel, \ + fp8_t> \ + <<>>( \ + out.mutable_data_ptr(), \ + input.const_data_ptr(), \ + scale.const_data_ptr(), d); \ + }); \ }); -void silu_and_mul_quant(torch::Tensor& out, // [..., d] - torch::Tensor& input, // [..., 2 * d] - torch::Tensor& scale) { - TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn || - out.dtype() == torch::kFloat8_e4m3fnuz); - TORCH_CHECK(input.dtype() == torch::kFloat16 || - input.dtype() == torch::kBFloat16); - TORCH_CHECK(input.size(-1) % 2 == 0); +void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + torch::stable::Tensor& scale) { + STD_TORCH_CHECK( + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn || + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz); + STD_TORCH_CHECK( + input.scalar_type() == torch::headeronly::ScalarType::Half || + input.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "Input must be FP16 or BF16"); + STD_TORCH_CHECK(input.size(-1) % 2 == 0); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel); } void persistent_masked_m_silu_mul_quant( - const at::Tensor& input, // (E, T, 2*H) - const at::Tensor& tokens_per_expert, // (E) - at::Tensor& y_q, // (E, T, H) [OUT] - at::Tensor& y_s, // (E, T, H//group_size) [OUT] + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] bool cast_scale_ue8m0) { #ifndef USE_ROCM @@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant( // fixed GROUP_SIZE of 128. static constexpr int GROUP_SIZE = 128; - TORCH_CHECK(input.dtype() == torch::kBFloat16); - TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn || - y_q.dtype() == torch::kFloat8_e4m3fnuz); - TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0); + STD_TORCH_CHECK(input.scalar_type() == + torch::headeronly::ScalarType::BFloat16); + STD_TORCH_CHECK( + y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn || + y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz); + STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0); bool const is_packed_ue8m0 = - (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0); - TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0); + (y_s.scalar_type() == torch::headeronly::ScalarType::Int && + cast_scale_ue8m0); + STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float || + is_packed_ue8m0); using Idx_t = int64_t; @@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant( int const NUM_GROUPS = H / GROUP_SIZE; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); // TODO: Get this from cuda_arch ? static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32; @@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant( static constexpr int max_shared_mem_bytes = \ GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \ dim3 grid(sms), block(THREAD_COUNT); \ - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ - VLLM_DISPATCH_FP8_TYPES( \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + VLLM_STABLE_DISPATCH_FP8_TYPES( \ y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \ vllm::silu_mul_fp8_quant_deep_gemm_kernel< \ BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \ Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \ <<>>( \ - reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \ - (fp8_t*)y_q.data_ptr(), \ - reinterpret_cast(y_s.data_ptr()), \ - reinterpret_cast(tokens_per_expert.data_ptr()), E, \ - T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \ + reinterpret_cast( \ + input.const_data_ptr()), \ + y_q.mutable_data_ptr(), \ + reinterpret_cast(y_s.mutable_data_ptr()), \ + reinterpret_cast( \ + tokens_per_expert.const_data_ptr()), \ + E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \ stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \ STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \ }); @@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant( Idx_t stride_ys_g = y_s.stride(2); Idx_t stride_ys_p = 0; if (!cast_scale_ue8m0) { - TORCH_CHECK(!is_packed_ue8m0); + STD_TORCH_CHECK(!is_packed_ue8m0); LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p, false); return; @@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant( return; } - TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0); - TORCH_CHECK(y_s.dtype() == torch::kInt32); + STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0); + STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int); // Int32 packed ue8m0 scales tensor. // Let E, T, G be the number to experts, number of tokens and number of groups diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh b/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh index 5cda4c9750a..72176179125 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh +++ b/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh @@ -100,6 +100,8 @@ void run_get_group_gemm_starts( int64_t k = a_tensors.size(1); int64_t scale_k = cutlass::ceil_div(k, b_group_size); + const torch::stable::accelerator::DeviceGuard device_guard( + a_tensors.get_device_index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index()); if (false) { diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu b/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu index 53ffe521363..1d7eb04093a 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu +++ b/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu @@ -17,11 +17,11 @@ #include #include #include "libtorch_stable/torch_utils.h" -#include "cutlass_extensions/torch_utils.hpp" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" #include "get_group_starts.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "w4a8_utils.cuh" namespace vllm::cutlass_w4a8_moe { diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu b/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu index 502f430b30b..8f19394b390 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu +++ b/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu @@ -6,7 +6,7 @@ #include #include #include "libtorch_stable/torch_utils.h" -#include "cutlass_extensions/torch_utils.hpp" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" #include "w4a8_utils.cuh" #include "cutlass/cutlass.h" @@ -22,7 +22,7 @@ #include "cutlass/util/mixed_dtype_utils.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu index 04e98b6076a..e4d2f220125 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -142,6 +142,8 @@ void mxfp4_run_get_group_gemm_starts( 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); + const torch::stable::accelerator::DeviceGuard device_guard( + a_tensors.get_device_index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index()); STD_TORCH_CHECK(out_tensors.size(1) == N, @@ -172,6 +174,8 @@ void run_mxfp4_blockwise_scaled_group_mm_sm100( const torch::stable::Tensor& problem_sizes, const torch::stable::Tensor& expert_offsets, const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + const torch::stable::accelerator::DeviceGuard device_guard( + a.get_device_index()); using ProblemShape = cutlass::gemm::GroupProblemShape>; using ElementType = cutlass::float_e2m1_t; diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 062f6018653..20f024bcef5 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -27,15 +27,24 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/cutlass_extensions/common.hpp" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" #include "nvfp4_utils.cuh" + +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12090 + #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 1 static_assert(CVT_FP4_ELTS_PER_THREAD == 16, "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)"); +#else + #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 0 +#endif #include "libtorch_stable/launch_bounds_utils.h" +#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED + namespace vllm { // MXFP4 block size constants @@ -104,7 +113,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) &input_offset_by_experts[chunk_start + 12])); local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]); -#pragma unroll + #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]; @@ -309,14 +318,14 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, } // 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); + /*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; @@ -364,12 +373,28 @@ static void validate_mxfp4_experts_quant_inputs( STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k); } +#endif // VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED + +static bool mxfp4_experts_quant_sm_supported(int64_t cuda_device_capability) { +#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED + return cuda_device_capability >= 100 && cuda_device_capability < 120; +#else + return false; +#endif +} + 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) { +#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED + int32_t sm = get_sm_version_num(); + STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), + "No compiled MXFP4 experts quant kernel for SM ", sm, + ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9."); + auto m_topk = input.size(0); auto k = input.size(1); @@ -390,6 +415,10 @@ void mxfp4_experts_quant( output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, stream); }); +#else + STD_TORCH_CHECK_NOT_IMPLEMENTED(false, + "MXFP4 experts quant requires CUDA >= 12.9."); +#endif } void silu_and_mul_mxfp4_experts_quant( @@ -398,6 +427,12 @@ void silu_and_mul_mxfp4_experts_quant( torch::stable::Tensor const& input_offset_by_experts, torch::stable::Tensor const& output_scale_offset_by_experts, int64_t n_experts) { +#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED + int32_t sm = get_sm_version_num(); + STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), + "No compiled SiLU+Mul MXFP4 experts quant kernel for SM ", sm, + ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9."); + 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)"); @@ -420,13 +455,29 @@ void silu_and_mul_mxfp4_experts_quant( output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, stream); }); +#else + STD_TORCH_CHECK_NOT_IMPLEMENTED( + false, "SiLU+Mul MXFP4 experts quant requires CUDA >= 12.9."); +#endif } -// 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. +bool mxfp4_experts_quant_supported(int64_t cuda_device_capability) { + return mxfp4_experts_quant_sm_supported(cuda_device_capability); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C, m) { + m.def("mxfp4_experts_quant_supported(int cuda_device_capability) -> bool"); +} + +// Registered here so the CUDA 12.8 stub and CUDA 12.9+ implementation stay +// tied to the same translation unit. 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)); } + +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) { + m.impl("mxfp4_experts_quant_supported", + TORCH_BOX(&mxfp4_experts_quant_supported)); +} diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu index 88caf03fda3..b044db6d32c 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu @@ -173,6 +173,8 @@ void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts, torch::stable::Tensor const& problem_sizes, int M, int N, int K) { int num_experts = (int)expert_offsets.size(0); + const torch::stable::accelerator::DeviceGuard device_guard( + a_tensors.get_device_index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index()); STD_TORCH_CHECK(out_tensors.size(1) == N, @@ -206,6 +208,8 @@ void run_fp4_blockwise_scaled_group_mm_sm100( const torch::stable::Tensor& problem_sizes, const torch::stable::Tensor& expert_offsets, const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + const torch::stable::accelerator::DeviceGuard device_guard( + a.get_device_index()); using ProblemShape = cutlass::gemm::GroupProblemShape>; using ElementType = cutlass::float_e2m1_t; @@ -411,6 +415,8 @@ void run_fp4_blockwise_scaled_group_mm_sm120( const torch::stable::Tensor& problem_sizes, const torch::stable::Tensor& expert_offsets, const torch::stable::Tensor& sf_offsets, int M, int N, int K) { + const torch::stable::accelerator::DeviceGuard device_guard( + a.get_device_index()); using ProblemShape = cutlass::gemm::GroupProblemShape>; using ElementType = cutlass::float_e2m1_t; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu index 86355bf7060..af9f24a70e0 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -31,7 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "core/batch_invariant.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 7adba6308fa..3a45ede8dfd 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -31,7 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "core/batch_invariant.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh index 0c04f010888..667138f3487 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh @@ -22,15 +22,15 @@ #include "../../cuda_vec_utils.cuh" -#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \ - CUDA_VERSION >= 12090 +#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDART_VERSION) && \ + CUDART_VERSION >= 12090 #define ELTS_PER_THREAD 16 + #define CVT_FP4_PACK16 1 constexpr int CVT_FP4_ELTS_PER_THREAD = 16; -constexpr bool CVT_FP4_PACK16 = true; #else #define ELTS_PER_THREAD 8 + #define CVT_FP4_PACK16 0 constexpr int CVT_FP4_ELTS_PER_THREAD = 8; -constexpr bool CVT_FP4_PACK16 = false; #endif constexpr int CVT_FP4_SF_VEC_SIZE = 16; @@ -237,21 +237,30 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4( // Get the final absolute maximum values. float vecMax = float(__hmax(localMax.x, localMax.y)); - // Get the SF (max value of the vector / max value of e2m1). - // maximum value of e2m1 = 6.0. - // TODO: use half as compute data type. - float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); // 8 bits representation of the SF. + float SFValue; uint8_t fp8SFVal; - // Write the SF to global memory (STG.8). + if constexpr (UE8M0_SF) { - // Extract the 8 exponent bits from float32. - // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits. - uint32_t tmp = reinterpret_cast(SFValue) >> 23; - fp8SFVal = tmp & 0xff; - // Convert back to fp32. - reinterpret_cast(SFValue) = tmp << 23; + // OCP MX spec E8M0 scale computation (MXFP4 path): + // scale_exp = biased_exponent(round_up(vecMax)) - 2 + // -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the + // safe divisor so that max_val / scale <= 6.0 for values near 2^n. + uint32_t max_bits = __float_as_uint(vecMax); + // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32 + // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n) + // round up to the next power of 2. + uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u; + uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu; + uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u; + scale_exp = min(scale_exp, 254u); + fp8SFVal = static_cast(scale_exp); + // Reconstruct scale as float32: scale = 2^(scale_exp - 127) + uint32_t sf_bits = scale_exp << 23; + SFValue = __uint_as_float(sf_bits); } else { + // NVFP4 path: scale = max / 6.0, stored as E4M3. + SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); // Here SFValue is always positive, so E4M3 is the same as UE4M3. __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp; @@ -262,13 +271,21 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4( // Write the SF to global memory (STG.8). if (SFout) *SFout = fp8SFVal; - // Get the output scale. - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * - // reciprocal(SFScaleVal)) - float outputScale = - SFValue != 0.0f ? reciprocal_approximate_ftz( + // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where + // SFScaleVal=1). Use exact division for UE8M0 to ensure bit-exact scaling + // that matches the reference QDQ implementation (dividing by a power-of-2 + // scale is exact in IEEE 754). + float outputScale; + if constexpr (UE8M0_SF) { + // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact. + outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f; + } else { + // NVFP4 path: use fast approximate reciprocal (original behavior). + outputScale = SFValue != 0.0f + ? reciprocal_approximate_ftz( SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; + } // Convert the input to float. float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu b/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu index bab7ac2a9f7..b32a7bd271f 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu +++ b/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu @@ -8,23 +8,11 @@ namespace vllm { -// Logic: one WARP per (token, group). Each of the 32 lanes owns EPT = -// group_size / 32 contiguous elements, sweeps the group in a single coalesced -// wide vector load (gate and up), and the per-group abs-max is a pure -// warp-shuffle reduction -- no shared memory and no __syncthreads. -// kWarpsPerBlock groups are packed into one 128-thread block. -// -// This replaces the previous one-block-per-group design, whose log2(group_size) -// shared-memory tree reduction (one __syncthreads per level) and one-element -// scalar loads left the kernel memory-latency-bound and ~2x off peak. The op -// sequence (fp32 SiLU, fmaxf abs-max) is unchanged, and fmaxf is -// order-invariant, so the per-group scale stays bit-identical. -constexpr int kWarpsPerBlock = 4; // 4 groups per 128-thread block +// Logic: one thread block per (token, group) pair template -__global__ void -__launch_bounds__(kWarpsPerBlock * 32) silu_and_mul_per_block_quant_kernel( +__global__ void silu_and_mul_per_block_quant_kernel( scalar_out_t* __restrict__ out, // Output: [num_tokens, hidden_size] in // FP8/INT8 float* __restrict__ scales, // Output: [num_tokens, hidden_size / @@ -36,98 +24,82 @@ __launch_bounds__(kWarpsPerBlock * 32) silu_and_mul_per_block_quant_kernel( ) { static_assert((group_size & (group_size - 1)) == 0, "group_size must be a power of 2 for correct reduction"); - static_assert(group_size % 32 == 0, - "group_size must be a multiple of the warp size"); - // Elements per thread: 2 for group_size=64, 4 for group_size=128. Each maps - // to a single vector load/store per lane (4B/8B in, 2B/4B out). - constexpr int EPT = group_size / 32; - int const tid = threadIdx.x; - int const warp_id = tid >> 5; - int const lane_id = tid & 31; + // Grid: (num_tokens, num_groups) int const token_idx = blockIdx.x; + int const group_idx = blockIdx.y; + int const tid = threadIdx.x; // tid in [0, group_size) int const num_tokens = gridDim.x; - // num_groups is no longer gridDim.y (we pack kWarpsPerBlock groups per - // block), so recover it from the compile-time group_size. - int const num_groups = hidden_size / group_size; - int const group_idx = blockIdx.y * kWarpsPerBlock + warp_id; - if (group_idx >= num_groups) return; // whole warp exits together (no sync) - // Input layout: [gate || up] concatenated along the last dimension. Each lane - // owns the EPT contiguous elements at group_start + lane_id * EPT, so the - // warp reads the whole group as one fully-coalesced wide load for gate and - // for up. + // Input layout: [gate || up] concatenated along last dimension int const input_stride = hidden_size * 2; int const group_start = group_idx * group_size; - int const lane_base = group_start + lane_id * EPT; + + // Pointers to this token's data scalar_t const* token_input_gate = - input + token_idx * input_stride + lane_base; + input + token_idx * input_stride + group_start; scalar_t const* token_input_up = token_input_gate + hidden_size; - scalar_out_t* token_output = out + token_idx * hidden_size + lane_base; + scalar_out_t* token_output = out + token_idx * hidden_size + group_start; // Scale pointer for this group + int const num_groups = gridDim.y; float* group_scale_ptr = is_scale_transposed ? scales + group_idx * num_tokens + token_idx : scales + token_idx * num_groups + group_idx; - // Step 1: one wide vector load per lane for gate and up, then SiLU(gate) * up - // in fp32. (group_start and hidden_size are both multiples of EPT because - // group_size = 32 * EPT divides hidden_size, so these loads are aligned.) - struct alignas(sizeof(scalar_t) * EPT) InVec { - scalar_t v[EPT]; - }; - InVec const gate_v = *reinterpret_cast(token_input_gate); - InVec const up_v = *reinterpret_cast(token_input_up); + // Shared memory for reduction (compile-time sized) + __shared__ float shared_max[group_size]; - float result[EPT]; // SiLU(gate) * up, kept in registers - float thread_max = 0.0f; + // Step 1: Each thread loads one element, computes SiLU, stores in register + float gate = static_cast(token_input_gate[tid]); + float up = static_cast(token_input_up[tid]); + + // Compute SiLU(gate) * up + float sigmoid_gate = 1.0f / (1.0f + expf(-gate)); + float silu_gate = gate * sigmoid_gate; + float result = silu_gate * up; // Keep in register + + // Step 2: Reduce to find group max + shared_max[tid] = fabsf(result); + __syncthreads(); + +// Power-of-2 reduction (group_size guaranteed to be power of 2) #pragma unroll - for (int k = 0; k < EPT; ++k) { - float gate = static_cast(gate_v.v[k]); - float up = static_cast(up_v.v[k]); - float sigmoid_gate = 1.0f / (1.0f + expf(-gate)); - float silu_gate = gate * sigmoid_gate; - result[k] = silu_gate * up; - thread_max = fmaxf(thread_max, fabsf(result[k])); + for (int stride = group_size / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + shared_max[tid] = fmaxf(shared_max[tid], shared_max[tid + stride]); + } + __syncthreads(); } - // Step 2: per-group abs-max via warp-shuffle. fmaxf is order-invariant, so - // the group max (and therefore the scale) is bit-identical to the tree - // reduction. -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - thread_max = - fmaxf(thread_max, __shfl_xor_sync(0xffffffffu, thread_max, offset)); - } + // Step 3: Compute scale (thread 0), broadcast via shared memory + if (tid == 0) { + float group_max = shared_max[0]; - // Step 3: compute the group scale in registers; lane 0 writes it to global. - float const group_max = thread_max; - float const quant_range = quant_type_max_v; - float group_scale = group_max / quant_range; + float const quant_range = quant_type_max_v; + float group_scale = group_max / quant_range; - // Apply scale upper bound if provided - if (scale_ub != nullptr) { - group_scale = fminf(group_scale, *scale_ub); - } + // Apply scale upper bound if provided + if (scale_ub != nullptr) { + group_scale = fminf(group_scale, *scale_ub); + } - // Use minimum safe scaling factor - group_scale = fmaxf(group_scale, min_scaling_factor::val()); + // Use minimum safe scaling factor + group_scale = fmaxf(group_scale, min_scaling_factor::val()); - if (lane_id == 0) { + // Store scale to global memory *group_scale_ptr = group_scale; - } - // Step 4: quantize the EPT owned elements and write them with one wide store. - struct alignas(sizeof(scalar_out_t) * EPT) OutVec { - scalar_out_t q[EPT]; - }; - OutVec out_v; -#pragma unroll - for (int k = 0; k < EPT; ++k) { - out_v.q[k] = vllm::ScaledQuant::quant_fn(result[k], - group_scale); + // Reuse shared_max[0] to broadcast scale + shared_max[0] = group_scale; } - *reinterpret_cast(token_output) = out_v; + __syncthreads(); + + float group_scale = shared_max[0]; + + // Step 4: Quantize and write output + token_output[tid] = + vllm::ScaledQuant::quant_fn(result, group_scale); } } // namespace vllm @@ -170,11 +142,8 @@ void silu_and_mul_per_block_quant(torch::stable::Tensor& out, input.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(input.get_device_index()); - // One warp per group; vllm::kWarpsPerBlock groups packed per 128-thread - // block. - dim3 grid(num_tokens, - (num_groups + vllm::kWarpsPerBlock - 1) / vllm::kWarpsPerBlock); - dim3 block(vllm::kWarpsPerBlock * 32); + dim3 grid(num_tokens, num_groups); + dim3 block(group_size); VLLM_STABLE_DISPATCH_FLOATING_TYPES( input.scalar_type(), "silu_and_mul_per_block_quant", [&] { diff --git a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh b/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh index 290abedcf94..0139cfbc956 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh @@ -8,7 +8,7 @@ #include "quantization/utils.cuh" #include "quant_conversions.cuh" -#include "../../../cub_helpers.h" +#include "../../cub_helpers.h" #include "../../../cuda_compat.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/gguf/dequantize.cuh b/csrc/libtorch_stable/quantization/gguf/dequantize.cuh deleted file mode 100644 index 9d355003ef9..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/dequantize.cuh +++ /dev/null @@ -1,571 +0,0 @@ -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/convert.cu -// Dequant functions -static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ - const block_q4_0 * x = (const block_q4_0 *) vx; - - const dfloat d = x[ib].d; - - const int vui = x[ib].qs[iqs]; - - v.x = __int2half_rn(vui & 0xF); - v.y = __int2half_rn(vui >> 4); - - v = __hsub2(v, __floats2half2_rn(8.0f, 8.0f)); - v = __hmul2(v, {d, d}); -} - -static __device__ __forceinline__ void dequantize_q4_1(const void * vx, const int ib, const int iqs, dfloat2 & v){ - const block_q4_1 * x = (const block_q4_1 *) vx; - - const dfloat d = __low2half(x[ib].dm); - const dfloat m = __high2half(x[ib].dm); - - const int vui = x[ib].qs[iqs]; - - v.x = __int2half_rn(vui & 0xF); - v.y = __int2half_rn(vui >> 4); - - v = __hmul2(v, {d, d}); - v = __hadd2(v, {m, m}); -} - -static __device__ __forceinline__ void dequantize_q5_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ - const block_q5_0 * x = (const block_q5_0 *) vx; - - const dfloat d = x[ib].d; - - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; - const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; - - v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); - v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); - - v = __hsub2(v, __floats2half2_rn(16.0f, 16.0f)); - v = __hmul2(v, {d, d}); -} - -static __device__ __forceinline__ void dequantize_q5_1(const void * vx, const int ib, const int iqs, dfloat2 & v){ - const block_q5_1 * x = (const block_q5_1 *) vx; - - const dfloat d = __low2half(x[ib].dm); - const dfloat m = __high2half(x[ib].dm); - - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; - const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; - - v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); - v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); - - v = __hmul2(v, {d, d}); - v = __hadd2(v, {m, m}); -} - -static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ - const block_q8_0 * x = (const block_q8_0 *) vx; - - const dfloat d = x[ib].d; - - v.x = __int2half_rn(x[ib].qs[iqs + 0]); - v.y = __int2half_rn(x[ib].qs[iqs + 1]); - - v = __hmul2(v, {d, d}); -} - -template -static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int k) { - const int i = 2*(blockDim.x*blockIdx.x + threadIdx.x); - - if (i >= k) { - return; - } - - const int ib = i/qk; // block index - const int iqs = (i%qk)/qr; // quant index - const int iybs = i - i%qk; // y block start index - const int y_offset = qr == 1 ? 1 : qk/2; - - // dequantize - dfloat2 v; - dequantize_kernel(vx, ib, iqs, v); - - y[iybs + iqs + 0] = convert_from_half(v.x); - y[iybs + iqs + y_offset] = convert_from_half(v.y); -} - -template -static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_q2_K * x = (const block_q2_K *) vx; - - const auto tid = threadIdx.x; - const int n = tid/32; - const int l = tid - 32*n; - const int is = 8*n + l/16; - - const uint8_t q = x[i].qs[32*n + l]; - dst_t * y = yy + i*QK_K + 128*n; - - half dall = __low2half(x[i].dm); - half dmin = __high2half(x[i].dm); - y[l+ 0] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+0] & 0xF) * ((q >> 0) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+0] >> 4)))); - y[l+32] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+2] & 0xF) * ((q >> 2) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+2] >> 4)))); - y[l+64] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+4] & 0xF) * ((q >> 4) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+4] >> 4)))); - y[l+96] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+6] & 0xF) * ((q >> 6) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+6] >> 4)))); -} - -template -static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_q3_K * x = (const block_q3_K *) vx; - - const auto r = threadIdx.x/4; - const int tid = r/2; - const int is0 = r%2; - const int l0 = 16*is0 + 4*(threadIdx.x%4); - const int n = tid / 4; - const int j = tid - 4*n; - - uint8_t m = 1 << (4*n + j); - int is = 8*n + 2*j + is0; - int shift = 2*j; - - int8_t us = is < 4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) : - is < 8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) : - is < 12 ? (x[i].scales[is-8] >> 4) | (((x[i].scales[is+0] >> 4) & 3) << 4) : - (x[i].scales[is-8] >> 4) | (((x[i].scales[is-4] >> 6) & 3) << 4); - half d_all = x[i].d; - half dl = __hmul(d_all, __int2half_rn(us - 32)); - - dst_t * y = yy + i*QK_K + 128*n + 32*j; - const uint8_t * q = x[i].qs + 32*n; - const uint8_t * hm = x[i].hmask; - - for (int l = l0; l < l0+4; ++l) { - y[l] = convert_from_half(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)))); - } -} - -static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) { - if (j < 4) { - d = q[j] & 63; m = q[j + 4] & 63; - } else { - d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); - m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); - } -} - -template -static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { - const block_q4_K * x = (const block_q4_K *) vx; - - const auto i = blockIdx.x; - - // assume 32 threads - const auto tid = threadIdx.x; - const int il = tid/8; - const int ir = tid%8; - const int is = 2*il; - const int n = 4; - - dst_t * y = yy + i*QK_K + 64*il + n*ir; - - const half dall = __low2half(x[i].dm); - const half dmin = __high2half(x[i].dm); - - const uint8_t * q = x[i].qs + 32*il + n*ir; - - uint8_t sc, m; - get_scale_min_k4(is + 0, x[i].scales, sc, m); - const half d1 = __hmul(dall, __int2half_rn(sc)); - const half m1 = __hmul(dmin, __int2half_rn(m)); - get_scale_min_k4(is + 1, x[i].scales, sc, m); - const half d2 = __hmul(dall, __int2half_rn(sc)); - const half m2 = __hmul(dmin, __int2half_rn(m)); - for (int l = 0; l < n; ++l) { - y[l + 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1)); - y[l +32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2)); - } -} - -template -static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { - const block_q5_K * x = (const block_q5_K *) vx; - - const auto i = blockIdx.x; - - // assume 64 threads - this is very slightly better than the one below - const auto tid = threadIdx.x; - const int il = tid/16; // il is in 0...3 - const int ir = tid%16; // ir is in 0...15 - const int is = 2*il; // is is in 0...6 - - dst_t * y = yy + i*QK_K + 64*il + 2*ir; - - const half dall = __low2half(x[i].dm); - const half dmin = __high2half(x[i].dm); - - const uint8_t * ql = x[i].qs + 32*il + 2*ir; - const uint8_t * qh = x[i].qh + 2*ir; - - uint8_t sc, m; - get_scale_min_k4(is + 0, x[i].scales, sc, m); - const half d1 = __hmul(dall, __int2half_rn(sc)); const half m1 = __hmul(dmin, __int2half_rn(m)); - get_scale_min_k4(is + 1, x[i].scales, sc, m); - const half d2 = __hmul(dall, __int2half_rn(sc)); const half m2 = __hmul(dmin, __int2half_rn(m)); - - uint8_t hm = 1 << (2*il); - y[ 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1)); - y[ 1] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1)); - hm <<= 1; - y[32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2)); - y[33] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2)); -} - -template -static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { - const block_q6_K * x = (const block_q6_K *) vx; - - const auto i = blockIdx.x; - - // assume 64 threads - this is very slightly better than the one below - const auto tid = threadIdx.x; - const int ip = tid/32; // ip is 0 or 1 - const int il = tid - 32*ip; // 0...32 - const int is = 8*ip + il/16; - - dst_t * y = yy + i*QK_K + 128*ip + il; - - const half d = x[i].d; - - const uint8_t * ql = x[i].ql + 64*ip + il; - const uint8_t qh = x[i].qh[32*ip + il]; - const int8_t * sc = x[i].scales + is; - - y[ 0] = convert_from_half(__hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)))); - y[32] = convert_from_half(__hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)))); - y[64] = convert_from_half(__hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)))); - y[96] = convert_from_half(__hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)))); -} - -template -static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq2_xxs * x = (const block_iq2_xxs *) vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint16_t * q2 = x[i].qs + 4*ib; - const uint8_t * aux8 = (const uint8_t *)q2; - const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]); - const uint32_t aux32 = q2[2] | (q2[3] << 16); - const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f; - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127]; - for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); -} - -template -static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq2_xs * x = (const block_iq2_xs *) vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint16_t * q2 = x[i].qs + 4*ib; - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511)); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; - const uint8_t signs = ksigns_iq2xs[q2[il] >> 9]; - for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); - -} - -template -static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq2_s * x = (const block_iq2_s *) vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300))); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; - const uint8_t signs = x[i].qs[QK_K/8+4*ib+il]; - for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); -} - -template -static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq3_xxs * x = (const block_iq3_xxs *) vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint8_t * q3 = x[i].qs + 8*ib; - const uint16_t * gas = (const uint16_t *)(x[i].qs + QK_K/4) + 2*ib; - const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]); - const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]); - const uint32_t aux32 = gas[0] | (gas[1] << 16); - const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f; - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127]; - for (int j = 0; j < 4; ++j) { - y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); - y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); - } -} - -template -static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq3_s * x = (const block_iq3_s *) vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint8_t * qs = x[i].qs + 8*ib; - const uint8_t * grid1 = (const uint8_t *)(iq3xs_grid + (qs[2*il+0] | ((x[i].qh[ib] << (8-2*il)) & 256))); - const uint8_t * grid2 = (const uint8_t *)(iq3xs_grid + (qs[2*il+1] | ((x[i].qh[ib] << (7-2*il)) & 256))); - const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf)) * 0.5f; - const uint8_t signs = x[i].signs[4*ib + il]; - for (int j = 0; j < 4; ++j) { - y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); - y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); - } -} - -template -static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const int64_t i = blockIdx.x; - const block_iq1_s * x = (const block_iq1_s *) vx; - - const int64_t tid = threadIdx.x; - const int64_t il = tid/8; // 0...3 - const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; - const float d = __half2float(x[i].d) * (2*((x[i].qh[ib] >> 12) & 7) + 1); - uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32; - grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[ib] >> 3*il) & 7) << 8)]; - grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; - grid32[0] &= 0x0f0f0f0f; - for (int j = 0; j < 8; ++j) { - y[j] = d * (q[j] + delta); - } -} - -template -static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const int64_t i = blockIdx.x; - const block_iq1_m * x = (const block_iq1_m *) vx; - - const int64_t tid = threadIdx.x; - const int64_t il = tid/8; // 0...3 - const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 8*il; - const uint16_t * sc = (const uint16_t *)x[i].scales; - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4); - const float d = __half2float(scale.f16) * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1); - const float delta = x[i].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA; - uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32; - grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)]; - grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; - grid32[0] &= 0x0f0f0f0f; - for (int j = 0; j < 8; ++j) { - y[j] = d * (q[j] + delta); - } -} - -template -static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) { - - const auto i = blockIdx.x; - const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL); - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 4*il; - const uint8_t * q4 = x[ib].qs + 4*il; - const float d = __half2float(x[ib].d); - for (int j = 0; j < 4; ++j) { - y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf]; - y[j+16] = d * kvalues_iq4nl[q4[j] >> 4]; - } - -} - -template -static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) { - const auto i = blockIdx.x; - const block_iq4_xs * x = (const block_iq4_xs *)vx; - - const auto tid = threadIdx.x; - const int il = tid/8; // 0...3 - const int ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 4*il; - const uint8_t * q4 = x[i].qs + 16*ib + 4*il; - const float d = __half2float(x[i].d) * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32); - for (int j = 0; j < 4; ++j) { - y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf]; - y[j+16] = d * kvalues_iq4nl[q4[j] >> 4]; - } -} - -template -static void dequantize_block_cuda(const void * __restrict__ vx, dst_t * __restrict__ y, const int k, cudaStream_t stream) { - const int num_blocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE); - dequantize_block<<>>(vx, y, k); -} - -template -static void dequantize_row_q2_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q2_K<<>>(vx, y); -} - -template -static void dequantize_row_q3_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q3_K<<>>(vx, y); -} - -template -static void dequantize_row_q4_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q4_K<<>>(vx, y); -} - -template -static void dequantize_row_q5_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q5_K<<>>(vx, y); -} - -template -static void dequantize_row_q6_K_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_q6_K<<>>(vx, y); -} - -template -static void dequantize_row_iq2_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_xxs<<>>(vx, y); -} - -template -static void dequantize_row_iq2_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_xs<<>>(vx, y); -} - -template -static void dequantize_row_iq2_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq2_s<<>>(vx, y); -} - -template -static void dequantize_row_iq3_xxs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq3_xxs<<>>(vx, y); -} - -template -static void dequantize_row_iq3_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq3_s<<>>(vx, y); -} - -template -static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq1_s<<>>(vx, y); -} - -template -static void dequantize_row_iq1_m_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = k / QK_K; - dequantize_block_iq1_m<<>>(vx, y); -} - -template -static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_iq4_nl<<>>(vx, y); -} - -template -static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int k, cudaStream_t stream) { - const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_iq4_xs<<>>(vx, y); -} - -template -static to_cuda_ggml_t ggml_get_to_cuda(int64_t type) { - switch (type) { - case 2: - return dequantize_block_cuda; - case 3: - return dequantize_block_cuda; - case 6: - return dequantize_block_cuda; - case 7: - return dequantize_block_cuda; - case 8: - return dequantize_block_cuda; - case 10: - return dequantize_row_q2_K_cuda; - case 11: - return dequantize_row_q3_K_cuda; - case 12: - return dequantize_row_q4_K_cuda; - case 13: - return dequantize_row_q5_K_cuda; - case 14: - return dequantize_row_q6_K_cuda; - case 16: - return dequantize_row_iq2_xxs_cuda; - case 17: - return dequantize_row_iq2_xs_cuda; - case 18: - return dequantize_row_iq3_xxs_cuda; - case 19: - return dequantize_row_iq1_s_cuda; - case 20: - return dequantize_row_iq4_nl_cuda; - case 21: - return dequantize_row_iq3_s_cuda; - case 22: - return dequantize_row_iq2_s_cuda; - case 23: - return dequantize_row_iq4_xs_cuda; - case 29: - return dequantize_row_iq1_m_cuda; - default: - return nullptr; - } -} diff --git a/csrc/libtorch_stable/quantization/gguf/ggml-common.h b/csrc/libtorch_stable/quantization/gguf/ggml-common.h deleted file mode 100644 index 6bef5db3ccf..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/ggml-common.h +++ /dev/null @@ -1,1150 +0,0 @@ -// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-common.h -#define QK_K 256 -#define K_QUANTS_PER_ITERATION 2 -#define WARP_SIZE_GGUF 32 -#define K_SCALE_SIZE 12 -#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 -#define CUDA_QUANTIZE_BLOCK_SIZE 256 -#define GGML_CUDA_DMMV_X 32 -#define GGML_CUDA_MMV_Y 1 - - -// Data Structures -// QK = number of values after dequantization -// QR = QK / number of values before dequantization -// QI = number of 32 bit integers before dequantization - -#define QK4_0 32 -#define QR4_0 2 -#define QI4_0 (QK4_0 / (4 * QR4_0)) -typedef struct { - half d; // delta - uint8_t qs[QK4_0 / 2]; // nibbles / quants -} block_q4_0; - -#define QK4_1 32 -#define QR4_1 2 -#define QI4_1 (QK4_1 / (4 * QR4_1)) -typedef struct { - half2 dm; // dm.x = delta, dm.y = min - uint8_t qs[QK4_1 / 2]; // nibbles / quants -} block_q4_1; - -#define QK5_0 32 -#define QR5_0 2 -#define QI5_0 (QK5_0 / (4 * QR5_0)) -typedef struct { - half d; // delta - uint8_t qh[4]; // 5-th bit of quants - uint8_t qs[QK5_0 / 2]; // nibbles / quants -} block_q5_0; - -#define QK5_1 32 -#define QR5_1 2 -#define QI5_1 (QK5_1 / (4 * QR5_1)) -typedef struct { - half2 dm; // dm.x = delta, dm.y = min - uint8_t qh[4]; // 5-th bit of quants - uint8_t qs[QK5_1 / 2]; // nibbles / quants -} block_q5_1; - -#define QK8_0 32 -#define QR8_0 1 -#define QI8_0 (QK8_0 / (4 * QR8_0)) -typedef struct { - half d; // delta - int8_t qs[QK8_0]; // quants -} block_q8_0; - -#define QK8_1 32 -#define QR8_1 1 -#define QI8_1 (QK8_1 / (4 * QR8_1)) -typedef struct { - half2 ds; // ds.x = delta, ds.y = sum - int8_t qs[QK8_0]; // quants -} block_q8_1; - -#define QR2_K 4 -#define QI2_K (QK_K / (4*QR2_K)) -typedef struct { - uint8_t scales[QK_K/16]; // scales and mins, quantized with 4 bits - uint8_t qs[QK_K/4]; // quants - half2 dm; // super-block scale for quantized scales/mins -} block_q2_K; - -#define QR3_K 4 -#define QI3_K (QK_K / (4*QR3_K)) -typedef struct { - uint8_t hmask[QK_K/8]; // quants - high bit - uint8_t qs[QK_K/4]; // quants - low 2 bits - uint8_t scales[K_SCALE_SIZE]; // scales, quantized with 6 bits - half d; // super-block scale -} block_q3_K; - -#define QR4_K 2 -#define QI4_K (QK_K / (4*QR4_K)) -typedef struct { - half2 dm; // super-block scale for quantized scales/mins - uint8_t scales[3*QK_K/64]; // scales, quantized with 6 bits - uint8_t qs[QK_K/2]; // 4--bit quants -} block_q4_K; - -#define QR5_K 2 -#define QI5_K (QK_K / (4*QR5_K)) -typedef struct { - half2 dm; // super-block scale for quantized scales/mins - uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits - uint8_t qh[QK_K/8]; // quants, high bit - uint8_t qs[QK_K/2]; // quants, low 4 bits -} block_q5_K; - -#define QR6_K 2 -#define QI6_K (QK_K / (4*QR6_K)) -typedef struct { - uint8_t ql[QK_K/2]; // quants, lower 4 bits - uint8_t qh[QK_K/4]; // quants, upper 2 bits - int8_t scales[QK_K/16]; // scales - half d; // delta -} block_q6_K; - -#define QR2_XXS 8 -#define QI2_XXS (QK_K / (4*QR2_XXS)) -typedef struct { - half d; - uint16_t qs[QK_K/8]; -} block_iq2_xxs; - -#define QR2_XS 8 -#define QI2_XS (QK_K / (4*QR2_XS)) -typedef struct { - half d; - uint16_t qs[QK_K/8]; - uint8_t scales[QK_K/32]; -} block_iq2_xs; - -#define QR2_S 8 -#define QI2_S (QK_K / (4*QR2_S)) -typedef struct { - half d; - uint8_t qs[QK_K/4]; - uint8_t qh[QK_K/32]; - uint8_t scales[QK_K/32]; -} block_iq2_s; - -#define QR3_XXS 8 -#define QI3_XXS (QK_K / (4*QR3_XXS)) -typedef struct { - half d; - uint8_t qs[3*(QK_K/8)]; -} block_iq3_xxs; - -#define QR3_XS 8 -#define QI3_XS (QK_K / (4*QR3_XS)) -#define IQ3S_N_SCALE QK_K/64 -typedef struct { - half d; - uint8_t qs[QK_K/4]; - uint8_t qh[QK_K/32]; - uint8_t signs[QK_K/8]; - uint8_t scales[IQ3S_N_SCALE]; -} block_iq3_s; - -// 1.5625 bpw -#define QR1_S 8 -#define QI1_S (QK_K / (4*QR1_S)) -typedef struct { - half d; - uint8_t qs[QK_K/8]; - uint16_t qh[QK_K/32]; -} block_iq1_s; - -// 1.75 bpw -#define QR1_M 8 -#define QI1_M (QK_K / (4*QR1_M)) -typedef struct { - uint8_t qs[QK_K/8]; // grid index, low 8 bits - uint8_t qh[QK_K/16]; // grid index, high 3 bits + grid shift bit (for two groups of 8) - uint8_t scales[QK_K/32]; // 3-bit block scales (4-bit if QK_K == 64) -} block_iq1_m; - -// Used by IQ1_M quants -typedef union { - half f16; - uint16_t u16; -} iq1m_scale_t; - -#define QK4_NL 32 -#define QR4_NL 2 -#define QI4_NL (QK4_NL / (4*QR4_NL)) -typedef struct { - half d; - uint8_t qs[QK4_NL/2]; -} block_iq4_nl; - -#define QR4_XS 8 -#define QI4_XS (QK_K / (4*QR4_XS)) -typedef struct { - half d; - uint16_t scales_h; - uint8_t scales_l[QK_K/64]; - uint8_t qs[QK_K/2]; -} block_iq4_xs; - -static const __device__ uint64_t iq2xxs_grid[256] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, - 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, - 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, - 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, - 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, - 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, - 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, - 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, - 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, - 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, - 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, - 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, - 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, - 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, - 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, - 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, - 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, - 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, - 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, - 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, - 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, - 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, - 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, - 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, - 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, - 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, - 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, - 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, - 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, - 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, - 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, - 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, - 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, - 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, - 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, - 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, - 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, - 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, - 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, - 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, - 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, - 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, - 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, - 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, - 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, - 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, - 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, - 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, - 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, - 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, - 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, - 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, - 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, - 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, - 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, - 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, - 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, - 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, - 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, - 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, - 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, - 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, - 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, - 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, -}; - -static const __device__ uint64_t iq2xs_grid[512] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, - 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, - 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, - 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, - 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, - 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, - 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, - 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, - 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, - 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, - 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, - 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, - 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, - 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, - 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, - 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, - 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, - 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, - 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, - 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, - 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, - 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, - 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, - 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, - 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, - 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, - 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, - 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, - 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, - 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, - 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, - 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, - 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, - 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, - 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, - 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, - 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, - 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, - 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, - 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, - 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, - 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, - 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, - 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, - 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, - 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, - 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, - 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, - 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, - 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, - 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, - 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, - 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, - 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, - 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, - 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, - 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, - 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, - 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, - 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, - 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, - 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, - 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, - 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, - 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, - 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, - 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, - 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, - 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, - 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, - 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, - 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, - 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, - 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, - 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, - 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, - 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, - 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, - 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, - 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, - 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, - 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, - 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, - 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, - 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, - 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, - 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, - 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, - 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, - 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, - 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, - 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, - 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, - 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, - 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, - 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, - 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, - 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, - 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, - 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, - 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, - 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, - 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, - 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, - 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, - 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, - 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, - 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, - 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, - 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, - 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, - 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, - 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, - 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, - 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, - 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, - 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, - 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, - 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, - 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, - 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, - 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, - 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, - 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, - 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, - 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, - 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, - 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, -}; - -static const __device__ uint64_t iq2s_grid[1024] = { - 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, - 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, - 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, - 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, - 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, - 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, - 0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, - 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808, - 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908, - 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, - 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, - 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, - 0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19, - 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819, - 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, - 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, - 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, - 0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908, - 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919, - 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, - 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, - 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, - 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b, - 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, - 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, - 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, - 0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, - 0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08, - 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, - 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, - 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, - 0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, - 0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b, - 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908, - 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, - 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, - 0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, - 0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08, - 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819, - 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, - 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, - 0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, - 0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b, - 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919, - 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, - 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, - 0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, - 0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919, - 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808, - 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, - 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, - 0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, - 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, - 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, - 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, - 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, - 0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, - 0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b, - 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819, - 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, - 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, - 0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, - 0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908, - 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b, - 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, - 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, - 0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, - 0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819, - 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819, - 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, - 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, - 0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, - 0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819, - 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808, - 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, - 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, - 0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, - 0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19, - 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08, - 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, - 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, - 0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, - 0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819, - 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908, - 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, - 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, - 0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, - 0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819, - 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908, - 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, - 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, - 0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, - 0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b, - 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, - 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, - 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, - 0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, - 0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908, - 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, - 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, - 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, - 0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, - 0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08, - 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08, - 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, - 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, - 0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, - 0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819, - 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908, - 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, - 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, - 0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, - 0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808, - 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819, - 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, - 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, - 0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, - 0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, - 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, - 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, - 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, - 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, - 0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19, - 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819, - 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, - 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, - 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, - 0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b, - 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b, - 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, - 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, - 0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, - 0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908, - 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08, - 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, - 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, - 0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, - 0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908, - 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b, - 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, - 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, - 0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, - 0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08, - 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819, - 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, - 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, - 0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, - 0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808, - 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808, - 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, - 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, - 0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, - 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808, - 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819, - 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, - 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, - 0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, - 0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908, - 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808, - 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, - 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, - 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, - 0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b, - 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, - 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, - 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, - 0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, - 0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908, - 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b, - 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, - 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, - 0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, - 0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919, - 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808, - 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, - 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, - 0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, - 0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808, - 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819, - 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, - 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, - 0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, - 0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808, - 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b, - 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, - 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, - 0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, - 0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808, - 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919, - 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, - 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, - 0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, - 0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808, - 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b, - 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, - 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, - 0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, - 0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808, - 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908, - 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, - 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, - 0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, - 0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908, - 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808, - 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, - 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, - 0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, - 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b, - 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808, - 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, - 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, - 0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, - 0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08, - 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908, - 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, - 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, - 0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, - 0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b, - 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808, - 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, - 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, - 0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, - 0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808, - 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808, - 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, - 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, - 0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, - 0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908, - 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08, - 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, - 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, - 0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, - 0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819, - 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908, - 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, - 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, - 0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, - 0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b, - 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908, - 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, - 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, - 0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, - 0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819, - 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808, - 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, - 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, - 0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, - 0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, - 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b, - 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, - 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, - 0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, - 0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819, - 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, - 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, - 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, -}; - -static const __device__ uint32_t iq3xxs_grid[256] = { - 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, - 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, - 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, - 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, - 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, - 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c, - 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34, - 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c, - 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, - 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, - 0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, - 0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, - 0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, - 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, - 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e, - 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24, - 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24, - 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, - 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, - 0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, - 0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, - 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, - 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, - 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c, - 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c, - 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14, - 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, - 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, - 0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, - 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, - 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, - 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, -}; - -static const __device__ uint32_t iq3xs_grid[512] = { - 0x04040404, 0x0404040c, 0x04040414, 0x0404042c, 0x0404043e, 0x04040c04, 0x04040c0c, 0x04040c14, - 0x04040c24, 0x04040c34, 0x04041404, 0x0404140c, 0x0404142c, 0x04041c1c, 0x04042404, 0x04042414, - 0x0404242c, 0x0404243e, 0x04042c0c, 0x04042c1c, 0x04043404, 0x04043414, 0x04043e0c, 0x04043e24, - 0x04043e3e, 0x040c0404, 0x040c040c, 0x040c0414, 0x040c0424, 0x040c0c04, 0x040c0c0c, 0x040c0c2c, - 0x040c1404, 0x040c141c, 0x040c143e, 0x040c1c0c, 0x040c1c2c, 0x040c2424, 0x040c340c, 0x040c342c, - 0x040c3e14, 0x04140404, 0x0414040c, 0x0414042c, 0x0414043e, 0x04140c04, 0x04140c1c, 0x04140c34, - 0x0414140c, 0x0414142c, 0x04141c04, 0x04141c24, 0x04142414, 0x0414242c, 0x0414243e, 0x04142c0c, - 0x04142c1c, 0x04143e04, 0x04143e1c, 0x041c041c, 0x041c0c0c, 0x041c0c2c, 0x041c1404, 0x041c1414, - 0x041c1c0c, 0x041c1c1c, 0x041c1c34, 0x041c2424, 0x041c2c04, 0x041c2c14, 0x041c343e, 0x041c3e0c, - 0x041c3e2c, 0x04240404, 0x04240c1c, 0x04240c3e, 0x0424140c, 0x04241424, 0x04241c14, 0x04242404, - 0x0424241c, 0x04242c0c, 0x04243e04, 0x042c0414, 0x042c0424, 0x042c1404, 0x042c1414, 0x042c1434, - 0x042c1c1c, 0x042c240c, 0x042c242c, 0x042c243e, 0x042c3434, 0x042c3e1c, 0x04340434, 0x04340c0c, - 0x04340c1c, 0x04341c0c, 0x04342c14, 0x04343e0c, 0x043e0404, 0x043e0414, 0x043e0424, 0x043e1404, - 0x043e1414, 0x043e1434, 0x043e1c1c, 0x043e2c04, 0x043e2c24, 0x0c040404, 0x0c04040c, 0x0c040414, - 0x0c040424, 0x0c040c04, 0x0c040c0c, 0x0c040c1c, 0x0c040c2c, 0x0c040c3e, 0x0c041404, 0x0c041414, - 0x0c041c0c, 0x0c041c24, 0x0c041c34, 0x0c042c24, 0x0c042c34, 0x0c04340c, 0x0c043e14, 0x0c0c0404, - 0x0c0c040c, 0x0c0c041c, 0x0c0c0434, 0x0c0c0c04, 0x0c0c0c24, 0x0c0c140c, 0x0c0c1c04, 0x0c0c1c1c, - 0x0c0c240c, 0x0c0c2c04, 0x0c0c2c14, 0x0c0c3e04, 0x0c0c3e34, 0x0c140404, 0x0c140c14, 0x0c140c2c, - 0x0c140c3e, 0x0c141404, 0x0c141424, 0x0c141c14, 0x0c142404, 0x0c14241c, 0x0c142c2c, 0x0c143404, - 0x0c143e14, 0x0c1c040c, 0x0c1c0424, 0x0c1c043e, 0x0c1c0c04, 0x0c1c0c1c, 0x0c1c140c, 0x0c1c143e, - 0x0c1c1c04, 0x0c1c1c24, 0x0c1c240c, 0x0c1c3414, 0x0c1c3e04, 0x0c24041c, 0x0c24042c, 0x0c240c14, - 0x0c240c24, 0x0c241c0c, 0x0c241c1c, 0x0c242414, 0x0c242434, 0x0c242c04, 0x0c242c24, 0x0c2c040c, - 0x0c2c0c04, 0x0c2c0c1c, 0x0c2c140c, 0x0c2c1c04, 0x0c2c1c14, 0x0c2c2c0c, 0x0c341404, 0x0c341424, - 0x0c34143e, 0x0c342424, 0x0c342434, 0x0c3e040c, 0x0c3e041c, 0x0c3e0c04, 0x0c3e0c14, 0x0c3e140c, - 0x0c3e1c2c, 0x0c3e240c, 0x0c3e3414, 0x0c3e3e04, 0x14040404, 0x1404040c, 0x1404041c, 0x1404042c, - 0x1404043e, 0x14040c04, 0x14040c14, 0x14040c24, 0x14040c34, 0x1404140c, 0x1404141c, 0x1404143e, - 0x14041c04, 0x14041c14, 0x1404240c, 0x1404241c, 0x1404242c, 0x14042c04, 0x14042c14, 0x1404343e, - 0x14043e04, 0x14043e1c, 0x14043e2c, 0x140c0404, 0x140c0414, 0x140c0c04, 0x140c0c1c, 0x140c0c3e, - 0x140c1414, 0x140c142c, 0x140c1c0c, 0x140c1c24, 0x140c2414, 0x140c2c0c, 0x1414040c, 0x14140424, - 0x1414043e, 0x1414140c, 0x1414141c, 0x14141c04, 0x14141c3e, 0x1414240c, 0x14142c1c, 0x14142c3e, - 0x14143e0c, 0x14143e24, 0x141c0404, 0x141c0414, 0x141c042c, 0x141c0c0c, 0x141c1414, 0x141c1424, - 0x141c1c0c, 0x141c1c1c, 0x141c2414, 0x141c2c04, 0x141c3434, 0x1424040c, 0x1424043e, 0x14241404, - 0x1424141c, 0x14241c14, 0x14241c2c, 0x1424240c, 0x14243e14, 0x14243e2c, 0x142c0424, 0x142c0c0c, - 0x142c1414, 0x142c1c3e, 0x142c2404, 0x142c2c1c, 0x142c3e04, 0x14340404, 0x14340414, 0x1434043e, - 0x1434140c, 0x14342c2c, 0x1434340c, 0x143e042c, 0x143e0c0c, 0x143e1434, 0x143e1c04, 0x143e241c, - 0x143e2c04, 0x1c040414, 0x1c040c0c, 0x1c040c1c, 0x1c040c2c, 0x1c040c3e, 0x1c041414, 0x1c041c0c, - 0x1c041c1c, 0x1c041c2c, 0x1c042414, 0x1c042424, 0x1c04243e, 0x1c042c0c, 0x1c04341c, 0x1c043e0c, - 0x1c0c040c, 0x1c0c041c, 0x1c0c042c, 0x1c0c0c24, 0x1c0c140c, 0x1c0c141c, 0x1c0c2404, 0x1c0c3404, - 0x1c0c3e14, 0x1c0c3e34, 0x1c140404, 0x1c140c14, 0x1c141404, 0x1c141c14, 0x1c141c24, 0x1c142c04, - 0x1c1c040c, 0x1c1c0c04, 0x1c1c0c24, 0x1c1c140c, 0x1c1c141c, 0x1c1c143e, 0x1c1c1c04, 0x1c1c240c, - 0x1c1c241c, 0x1c1c243e, 0x1c1c2c2c, 0x1c1c3e1c, 0x1c24041c, 0x1c240c0c, 0x1c240c34, 0x1c241414, - 0x1c241c0c, 0x1c242c14, 0x1c243404, 0x1c243424, 0x1c2c040c, 0x1c2c0c04, 0x1c2c0c14, 0x1c2c142c, - 0x1c2c1c14, 0x1c2c2424, 0x1c2c2c34, 0x1c2c3e1c, 0x1c340c34, 0x1c34240c, 0x1c3e040c, 0x1c3e041c, - 0x1c3e1404, 0x1c3e1414, 0x1c3e1c2c, 0x24040404, 0x24040424, 0x24040c14, 0x24041404, 0x24041424, - 0x2404143e, 0x24041c14, 0x2404240c, 0x24042c04, 0x24043e04, 0x240c0414, 0x240c043e, 0x240c0c0c, - 0x240c0c1c, 0x240c1414, 0x240c1c04, 0x240c1c2c, 0x240c241c, 0x240c2c0c, 0x240c2c2c, 0x2414040c, - 0x2414041c, 0x24140c04, 0x24140c2c, 0x2414140c, 0x24141c1c, 0x24142404, 0x24142c3e, 0x24143414, - 0x24143e04, 0x241c0424, 0x241c0c0c, 0x241c0c1c, 0x241c1404, 0x241c1414, 0x241c1c0c, 0x241c1c2c, - 0x24240404, 0x24240414, 0x24241424, 0x24241c3e, 0x24242404, 0x24243e0c, 0x242c042c, 0x242c043e, - 0x242c140c, 0x242c3414, 0x24340c1c, 0x24341c24, 0x24343404, 0x243e0c04, 0x243e0c2c, 0x243e1c04, - 0x243e241c, 0x243e2c0c, 0x2c040414, 0x2c040c04, 0x2c040c24, 0x2c041414, 0x2c042404, 0x2c042424, - 0x2c04243e, 0x2c042c14, 0x2c043434, 0x2c043e24, 0x2c0c040c, 0x2c0c041c, 0x2c0c042c, 0x2c0c0c14, - 0x2c0c140c, 0x2c0c1c14, 0x2c0c3e14, 0x2c140404, 0x2c140c0c, 0x2c14141c, 0x2c141c04, 0x2c141c34, - 0x2c142c1c, 0x2c1c0414, 0x2c1c043e, 0x2c1c0c04, 0x2c1c143e, 0x2c1c2424, 0x2c1c2c0c, 0x2c1c342c, - 0x2c1c3e1c, 0x2c24040c, 0x2c240424, 0x2c241404, 0x2c241c14, 0x2c242434, 0x2c2c0c14, 0x2c2c1434, - 0x2c2c2c0c, 0x2c2c2c1c, 0x2c342414, 0x2c3e0414, 0x2c3e0424, 0x2c3e1414, 0x34040c0c, 0x34040c1c, - 0x34040c2c, 0x34041c0c, 0x34041c1c, 0x34043404, 0x340c0404, 0x340c1404, 0x340c143e, 0x340c3424, - 0x34140c14, 0x34141c24, 0x34142414, 0x34142c2c, 0x34143414, 0x34143e04, 0x341c0404, 0x341c0c24, - 0x341c140c, 0x341c2404, 0x3424142c, 0x3424241c, 0x34243414, 0x342c0404, 0x342c041c, 0x342c1c24, - 0x342c3404, 0x3434042c, 0x34342404, 0x343e0c0c, 0x343e0c1c, 0x3e040404, 0x3e040424, 0x3e04043e, - 0x3e041404, 0x3e041414, 0x3e041c34, 0x3e042404, 0x3e042c24, 0x3e043414, 0x3e0c0414, 0x3e0c0c0c, - 0x3e0c1424, 0x3e0c241c, 0x3e0c242c, 0x3e14040c, 0x3e140424, 0x3e140c04, 0x3e140c34, 0x3e14140c, - 0x3e141c04, 0x3e142c0c, 0x3e1c0414, 0x3e1c1c14, 0x3e1c1c2c, 0x3e1c2c1c, 0x3e24040c, 0x3e24042c, - 0x3e240c1c, 0x3e241404, 0x3e242c04, 0x3e2c1414, 0x3e2c2414, 0x3e340414, 0x3e341c0c, 0x3e3e0404, -}; - -#define IQ1S_DELTA 0.125f -#define IQ1M_DELTA 0.125f -static const __device__ uint64_t iq1s_grid_gpu[2048] = { - 0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, - 0x00020002, 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, - 0x02000000, 0x02000002, 0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200, - 0x02020202, 0x00000110, 0x00000111, 0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212, - 0x00020111, 0x01000011, 0x01000112, 0x01000211, 0x01010012, 0x01010111, 0x01010212, 0x01020011, - 0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011, 0x02010110, 0x02010112, 0x02020111, - 0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020, 0x00020022, 0x00020220, - 0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020, 0x02000022, - 0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220, - 0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101, - 0x01011202, 0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110, - 0x00001111, 0x00001112, 0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111, - 0x01001212, 0x01011010, 0x01011011, 0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010, - 0x01021012, 0x01021111, 0x01021210, 0x01021212, 0x02001011, 0x02011011, 0x02011111, 0x02011210, - 0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112, 0x02021211, 0x00011120, 0x00011221, - 0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220, 0x01021020, 0x01021021, - 0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000, 0x00002002, - 0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101, - 0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101, - 0x02022000, 0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211, - 0x00022110, 0x00022111, 0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110, - 0x01022211, 0x02012011, 0x02012110, 0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022, - 0x00002220, 0x00002222, 0x00012121, 0x00022020, 0x00022022, 0x00022220, 0x00022222, 0x01002121, - 0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020, 0x02002022, 0x02002121, 0x02002220, - 0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222, 0x00110000, 0x00110001, - 0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000, 0x01110101, - 0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102, - 0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012, - 0x00110111, 0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010, - 0x01110011, 0x01110012, 0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111, - 0x02100110, 0x02110012, 0x02110111, 0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122, - 0x00120121, 0x01100020, 0x01100122, 0x01100221, 0x01110022, 0x01110121, 0x01110220, 0x01110222, - 0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120, 0x02110122, 0x02120121, 0x00101001, - 0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201, 0x00121001, 0x00121102, - 0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100, 0x01111101, - 0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000, - 0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101, - 0x02121201, 0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112, - 0x00111211, 0x00121010, 0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110, - 0x01101111, 0x01101112, 0x01111011, 0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211, - 0x01111212, 0x01121011, 0x01121110, 0x01121111, 0x01121112, 0x01121211, 0x02101010, 0x02101012, - 0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010, 0x02111011, 0x02111110, 0x02111111, - 0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111, 0x00101021, 0x00101120, - 0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021, 0x00121122, - 0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121, - 0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221, - 0x01121222, 0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001, - 0x00112102, 0x00122101, 0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101, - 0x01112200, 0x01112202, 0x01122000, 0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101, - 0x02112001, 0x02112100, 0x02122101, 0x00112010, 0x00112012, 0x00112111, 0x00112212, 0x00122011, - 0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210, 0x01112011, 0x01112110, 0x01112111, - 0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212, 0x02102211, 0x02112011, - 0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221, 0x00112122, - 0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121, - 0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222, - 0x00200000, 0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101, - 0x00220200, 0x00220202, 0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000, - 0x02200002, 0x02200200, 0x02200202, 0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200, - 0x02220202, 0x00200111, 0x00210011, 0x00210110, 0x00210211, 0x00220111, 0x01200012, 0x01200110, - 0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011, 0x01220110, 0x01220111, 0x01220112, - 0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021, 0x00200220, 0x00200222, - 0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121, 0x01210021, - 0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121, - 0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201, - 0x00221101, 0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200, - 0x01211202, 0x01221102, 0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101, - 0x00201211, 0x00211111, 0x00221011, 0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011, - 0x01211110, 0x01211111, 0x01211211, 0x01221012, 0x01221111, 0x01221210, 0x02201211, 0x02211010, - 0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011, 0x02221110, 0x02221112, 0x02221211, - 0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021, 0x01201221, 0x01211121, - 0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222, 0x00202000, - 0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202, - 0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202, - 0x02222000, 0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211, - 0x00222111, 0x01202112, 0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112, - 0x01222211, 0x02202111, 0x02212010, 0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020, - 0x00202022, 0x00202220, 0x00202222, 0x00222020, 0x00222022, 0x00222220, 0x00222222, 0x01202121, - 0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020, 0x02202022, 0x02202220, 0x02202222, - 0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101, 0x10010001, 0x10010102, - 0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001, 0x11020100, - 0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110, - 0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011, - 0x10020112, 0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111, - 0x11010112, 0x11010211, 0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110, - 0x12000112, 0x12010010, 0x12010012, 0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121, - 0x10010021, 0x10010120, 0x10010122, 0x10020121, 0x11000021, 0x11010022, 0x11010121, 0x11010222, - 0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121, 0x10001001, 0x10011101, 0x10011201, - 0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100, 0x11011101, 0x11011102, - 0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102, 0x12001201, - 0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012, - 0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010, - 0x10021111, 0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010, - 0x11011011, 0x11011110, 0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110, - 0x11021111, 0x11021112, 0x11021211, 0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011, - 0x12011110, 0x12011111, 0x12011112, 0x12011211, 0x12011212, 0x12021111, 0x12021210, 0x12021212, - 0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121, 0x10011220, 0x10011222, 0x10021021, - 0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220, 0x11011020, 0x11011021, - 0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220, 0x12001021, - 0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101, - 0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101, - 0x11012200, 0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100, - 0x12012102, 0x12012201, 0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010, - 0x10012110, 0x10012111, 0x10012210, 0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111, - 0x11002212, 0x11012011, 0x11012012, 0x11012110, 0x11012111, 0x11012112, 0x11012211, 0x11022010, - 0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112, 0x12002211, 0x12012012, 0x12012111, - 0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211, 0x10012122, 0x11002120, - 0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221, 0x12012120, - 0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101, - 0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001, - 0x11110100, 0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201, - 0x12110101, 0x12110200, 0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210, - 0x10100211, 0x10100212, 0x10110011, 0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211, - 0x10120010, 0x10120111, 0x10120112, 0x10120210, 0x10120212, 0x11100011, 0x11100110, 0x11100111, - 0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012, 0x11110110, 0x11110111, 0x11110112, - 0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111, 0x11120112, 0x11120211, - 0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211, 0x12120010, - 0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021, - 0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122, - 0x11110221, 0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221, - 0x12110222, 0x12120120, 0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102, - 0x10111200, 0x10111201, 0x10121001, 0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100, - 0x11101101, 0x11101102, 0x11101201, 0x11101202, 0x11111000, 0x11111001, 0x11111100, 0x11111101, - 0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001, 0x11121002, 0x11121100, 0x11121101, - 0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001, 0x12111100, 0x12111101, - 0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011, 0x10101012, - 0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110, - 0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112, - 0x10121211, 0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210, - 0x11101211, 0x11111010, 0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210, - 0x11111211, 0x11111212, 0x11121010, 0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210, - 0x11121211, 0x11121212, 0x12101011, 0x12101110, 0x12101111, 0x12101211, 0x12101212, 0x12111010, - 0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210, 0x12111211, 0x12121011, 0x12121110, - 0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022, 0x10101120, 0x10101122, - 0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221, 0x10121020, - 0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021, - 0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022, - 0x11111120, 0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120, - 0x11121121, 0x11121221, 0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222, - 0x12111021, 0x12111121, 0x12111222, 0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221, - 0x10102100, 0x10102101, 0x10102102, 0x10102201, 0x10112000, 0x10112101, 0x10112200, 0x10122001, - 0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001, 0x11112100, 0x11112101, 0x11112102, - 0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101, 0x12102002, 0x12102201, - 0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011, 0x10102012, - 0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111, - 0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012, - 0x11112110, 0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110, - 0x11122111, 0x11122112, 0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110, - 0x12112111, 0x12112112, 0x12112210, 0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121, - 0x10112222, 0x10122020, 0x10122121, 0x10122122, 0x10122221, 0x11102121, 0x11102220, 0x11102221, - 0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221, 0x11122022, 0x11122121, 0x11122220, - 0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122, 0x12112220, 0x12112222, - 0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100, 0x11210000, - 0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201, - 0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012, - 0x10210111, 0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011, - 0x11210111, 0x11210112, 0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212, - 0x12210012, 0x12210111, 0x12220011, 0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221, - 0x11200020, 0x11200021, 0x11200122, 0x11210121, 0x11210122, 0x11210220, 0x11220020, 0x12200121, - 0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002, 0x10211101, 0x10211102, 0x10211202, - 0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101, 0x11201200, 0x11201202, - 0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000, 0x11221002, - 0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101, - 0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210, - 0x10201212, 0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112, - 0x11201211, 0x11211010, 0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011, - 0x11221110, 0x11221111, 0x11221112, 0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011, - 0x12211111, 0x12211112, 0x12211211, 0x12211212, 0x12221012, 0x12221111, 0x12221112, 0x12221210, - 0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122, 0x10221220, 0x10221221, 0x11201020, - 0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121, 0x11211122, 0x11211220, - 0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121, 0x12201222, - 0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222, - 0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001, - 0x11222201, 0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010, - 0x10212111, 0x10222011, 0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111, - 0x11202112, 0x11202210, 0x11212011, 0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010, - 0x11222111, 0x11222212, 0x12202012, 0x12202110, 0x12202212, 0x12212111, 0x12222011, 0x12222110, - 0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220, 0x11202021, 0x11202120, 0x11202221, - 0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121, 0x11222221, 0x12202122, - 0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200, 0x20000202, - 0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100, - 0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101, - 0x22020000, 0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112, - 0x20010211, 0x20020111, 0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111, - 0x21010112, 0x21010210, 0x21010211, 0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211, - 0x22010110, 0x22010112, 0x22010211, 0x22020111, 0x20000020, 0x20000022, 0x20000220, 0x20000222, - 0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222, 0x21010021, 0x21010120, 0x21010221, - 0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121, 0x22020020, 0x22020022, - 0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001, 0x21011101, - 0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211, - 0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111, - 0x21001210, 0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111, - 0x21021112, 0x21021210, 0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010, - 0x22011012, 0x22011111, 0x22011210, 0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121, - 0x21001021, 0x21001120, 0x21001221, 0x21001222, 0x21011020, 0x21011121, 0x21011221, 0x21011222, - 0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021, 0x22011222, 0x22021120, 0x20002000, - 0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002, 0x20022200, 0x20022202, - 0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201, 0x22002000, - 0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202, - 0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110, - 0x21002112, 0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110, - 0x22002111, 0x22012112, 0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222, - 0x20012121, 0x20022020, 0x20022022, 0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120, - 0x21012122, 0x22002020, 0x22002022, 0x22002220, 0x22002222, 0x22012121, 0x22022020, 0x22022022, - 0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102, 0x20110200, 0x20110201, 0x20120101, - 0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202, 0x21120201, 0x21120202, - 0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011, 0x20100110, - 0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110, - 0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111, - 0x21110112, 0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111, - 0x22110210, 0x22120011, 0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120, - 0x20110221, 0x20120121, 0x21100120, 0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121, - 0x21110220, 0x21120122, 0x21120221, 0x22100121, 0x22110120, 0x22110122, 0x22120221, 0x20101001, - 0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200, 0x20121102, 0x21101000, 0x21101202, - 0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201, 0x21121000, 0x21121001, - 0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101, 0x22111200, - 0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011, - 0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212, - 0x21101011, 0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012, - 0x21111110, 0x21111111, 0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110, - 0x21121111, 0x21121112, 0x21121211, 0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012, - 0x22111110, 0x22111111, 0x22111112, 0x22111211, 0x22111212, 0x22121010, 0x22121012, 0x22121111, - 0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020, 0x20111121, 0x20111221, 0x20121020, - 0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021, 0x21111022, 0x21111121, - 0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221, 0x22101222, - 0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102, - 0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102, - 0x21112202, 0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101, - 0x20102110, 0x20102112, 0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212, - 0x20122010, 0x20122011, 0x20122110, 0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210, - 0x21102212, 0x21112011, 0x21112110, 0x21112111, 0x21112112, 0x21112211, 0x21122012, 0x21122111, - 0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010, 0x22112012, 0x22112111, 0x22112212, - 0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120, 0x21102122, 0x21102221, - 0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120, 0x22112121, - 0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002, - 0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000, - 0x22200002, 0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202, - 0x20200111, 0x20200211, 0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112, - 0x21200211, 0x21210011, 0x21210111, 0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111, - 0x22210010, 0x22210012, 0x22210112, 0x22210211, 0x20200022, 0x20200220, 0x20200222, 0x20210020, - 0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121, 0x21210021, 0x21210122, 0x21210221, - 0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121, 0x22220020, 0x22220022, - 0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000, 0x21211100, - 0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201, - 0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112, - 0x20221211, 0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211, - 0x21221111, 0x21221212, 0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012, - 0x22211111, 0x22211210, 0x20201121, 0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121, - 0x21201120, 0x21201122, 0x21201222, 0x21211022, 0x21211121, 0x21211122, 0x21211220, 0x21221020, - 0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122, 0x22211221, 0x22221021, 0x22221120, - 0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000, 0x20222002, 0x20222200, - 0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002, 0x22202200, - 0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110, - 0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011, - 0x21222112, 0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222, - 0x20222020, 0x20222022, 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, - 0x22202022, 0x22202220, 0x22202222, 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222, -}; - -static const __device__ uint8_t ksigns_iq2xs[128] = { - 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, - 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, - 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, - 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, - 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, - 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, - 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, - 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, -}; - -static const __device__ uint64_t ksigns64[128] = { - 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, - 0xff00000000ff0000, 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, - 0xff000000ff000000, 0x00000000ff0000ff, 0x00000000ff00ff00, 0xff000000ff00ffff, - 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, 0x00000000ffffffff, - 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff, - 0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, - 0x000000ffff000000, 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, - 0xff0000ffffff0000, 0x000000ffffff00ff, 0x000000ffffffff00, 0xff0000ffffffffff, - 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, 0xff00ff000000ffff, - 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff, - 0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, - 0xff00ff00ffff0000, 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, - 0x0000ffff00000000, 0xff00ffff000000ff, 0xff00ffff0000ff00, 0x0000ffff0000ffff, - 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, 0xff00ffff00ffffff, - 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff, - 0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, - 0xffff000000000000, 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, - 0x00ff000000ff0000, 0xffff000000ff00ff, 0xffff000000ffff00, 0x00ff000000ffffff, - 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, 0x00ff0000ff00ffff, - 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff, - 0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, - 0xffff00ff00ff0000, 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, - 0xffff00ffff000000, 0x00ff00ffff0000ff, 0x00ff00ffff00ff00, 0xffff00ffff00ffff, - 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, 0x00ff00ffffffffff, - 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff, - 0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, - 0xffffff00ff000000, 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, - 0x00ffff00ffff0000, 0xffffff00ffff00ff, 0xffffff00ffffff00, 0x00ffff00ffffffff, - 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, 0xffffffff0000ffff, - 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, - 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, - 0xffffffffffff0000, 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, -}; - -static const __device__ uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128}; -static const __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; - - -typedef half dfloat; // dequantize float -typedef half2 dfloat2; -typedef void (*dequantize_kernel_t)(const void * vx, const int ib, const int iqs, dfloat2 & v); -template -using to_cuda_ggml_t = void (*)(const void * __restrict__ x, dst_t * __restrict__ y, int k, cudaStream_t stream); -typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs); -typedef void (*allocate_tiles_cuda_t)(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc); -typedef void (*load_tiles_cuda_t)( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row); -typedef float (*vec_dot_q_mul_mat_cuda_t)( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ms, const int & i, const int & j, const int & k); - -// Utility function - -template -static __device__ __forceinline__ dst_t convert_from_half(half val) { - return val; -} - -template<> -__device__ __forceinline__ c10::BFloat16 convert_from_half(half val) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - return __float2bfloat16(__half2float(val)); -#else - return __half2float(val); -#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 -} - -template<> -__device__ __forceinline__ float convert_from_half(half val) { - return __half2float(val); -} - -#if defined(USE_ROCM) - -#ifndef __has_builtin - #define __has_builtin(x) 0 -#endif - -typedef int8_t int8x4_t __attribute__((ext_vector_type(4))); -static __device__ __forceinline__ int __vsubss4(const int a, const int b) { - const int8x4_t va = reinterpret_cast(a); - const int8x4_t vb = reinterpret_cast(b); -#if __has_builtin(__builtin_elementwise_sub_sat) - const int8x4_t c = __builtin_elementwise_sub_sat(va, vb); - return reinterpret_cast(c); -#else - int8x4_t c; - int16_t tmp; -#pragma unroll - for (int i = 0; i < 4; i++) { - tmp = va[i] - vb[i]; - if(tmp > std::numeric_limits::max()) tmp = std::numeric_limits::max(); - if(tmp < std::numeric_limits::min()) tmp = std::numeric_limits::min(); - c[i] = tmp; - } - return reinterpret_cast(c); -#endif // __has_builtin(__builtin_elementwise_sub_sat) -} - -static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) { -#if __has_builtin(__builtin_amdgcn_sdot4) - c = __builtin_amdgcn_sdot4(a, b, c, false); -#else - const int8x4_t va = reinterpret_cast(a); - const int8x4_t vb = reinterpret_cast(b); - c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3]; -#endif - return c; -} - -static __device__ __forceinline__ uint32_t __vcmpeq4(const uint32_t a, const uint32_t b) { - uint32_t neq = a^b; - return !(neq & 0xff000000) * 0xff000000 | - !(neq & 0x00ff0000) * 0x00ff0000 | - !(neq & 0x0000ff00) * 0x0000ff00 | - !(neq & 0x000000ff) * 0x000000ff; -} - -static __device__ __forceinline__ uint32_t __vsub4(const uint32_t a, const uint32_t b) { - return (static_cast(((a & 0xff000000) >> 24) - ((b & 0xff000000) >> 24)) << 24) + - (static_cast(((a & 0x00ff0000) >> 16) - ((b & 0x00ff0000) >> 16)) << 16) + - (static_cast(((a & 0x0000ff00) >> 8) - ((b & 0x0000ff00) >> 8)) << 8) + - (static_cast(((a & 0x000000ff) >> 0) - ((b & 0x000000ff) >> 0)) << 0); -} -#endif // defined(USE_ROCM) diff --git a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu b/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu deleted file mode 100644 index 2a56d7a18f4..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu +++ /dev/null @@ -1,557 +0,0 @@ -#include -#include - -#include "../../../cuda_compat.h" -#include "../../dispatch_utils.h" -#include "../../torch_utils.h" - -#include - -#include "ggml-common.h" -#include "vecdotq.cuh" -#include "dequantize.cuh" -#include "mmvq.cuh" -#include "mmq.cuh" -#include "moe.cuh" -#include "moe_vec.cuh" - -// Q8 gemv -template -static __global__ void quantize_q8_1(const scalar_t* __restrict__ x, - void* __restrict__ vy, const int kx, - const int kx_padded) { - const auto ix = blockDim.x * blockIdx.x + threadIdx.x; - if (ix >= kx_padded) { - return; - } - const auto iy = blockDim.y * blockIdx.y + threadIdx.y; - const int i_padded = iy * kx_padded + ix; - - block_q8_1* y = (block_q8_1*)vy; - - const int ib = i_padded / QK8_1; // block index - const int iqs = i_padded % QK8_1; // quant index - - const float xi = ix < kx ? static_cast(x[iy * kx + ix]) : 0.0f; - float amax = fabsf(xi); - float sum = xi; - -#pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) { - amax = fmaxf(amax, VLLM_SHFL_XOR_SYNC_WIDTH(amax, mask, 32)); - sum += VLLM_SHFL_XOR_SYNC_WIDTH(sum, mask, 32); - } - - const float d = amax / 127; - const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); - - y[ib].qs[iqs] = q; - - if (iqs > 0) { - return; - } - - y[ib].ds.x = __float2half(d); - y[ib].ds.y = __float2half(sum); -} - -template -static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx, - const int ky, cudaStream_t stream) { - const int64_t kx_padded = (kx + 512 - 1) / 512 * 512; - const int block_num_x = - (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; - constexpr int MAX_BLOCK_SIZE = 65535; - for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) { - const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off; - const dim3 num_blocks(block_num_x, num_blocks_y, 1); - const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1); - quantize_q8_1<<>>( - &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded); - } -} - -torch::stable::Tensor ggml_dequantize( - torch::stable::Tensor W, // quant weight - int64_t type, int64_t m, int64_t n, - std::optional const& dtype) { - const torch::stable::accelerator::DeviceGuard device_guard( - W.get_device_index()); - auto dtype_ = dtype.value_or(torch::headeronly::ScalarType::Half); - auto DW = torch::stable::empty({m, n}, dtype_, std::nullopt, W.device()); - cudaStream_t stream = get_current_cuda_stream(); - - VLLM_STABLE_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { - auto to_cuda = ggml_get_to_cuda(type); - to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); - }); - - return DW; -} - -torch::stable::Tensor ggml_mul_mat_vec_a8( - torch::stable::Tensor W, // quant weight - torch::stable::Tensor X, // input - int64_t type, int64_t row) { - int col = X.sizes()[1]; - int vecs = X.sizes()[0]; - const int padded = (col + 512 - 1) / 512 * 512; - const torch::stable::accelerator::DeviceGuard device_guard( - X.get_device_index()); - auto Y = torch::stable::empty({vecs, row}, X.scalar_type(), std::nullopt, - W.device()); - cudaStream_t stream = get_current_cuda_stream(); - auto quant_X = torch::stable::empty({vecs, padded / 32 * 9}, - torch::headeronly::ScalarType::Int, - std::nullopt, W.device()); - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - X.scalar_type(), "ggml_mul_mat_vec_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), - (void*)quant_X.data_ptr(), col, vecs, - stream); - switch (type) { - case 2: - mul_mat_vec_q4_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 3: - mul_mat_vec_q4_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 6: - mul_mat_vec_q5_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 7: - mul_mat_vec_q5_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 8: - mul_mat_vec_q8_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 10: - mul_mat_vec_q2_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 11: - mul_mat_vec_q3_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 12: - mul_mat_vec_q4_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 13: - mul_mat_vec_q5_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 14: - mul_mat_vec_q6_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 16: - mul_mat_vec_iq2_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 17: - mul_mat_vec_iq2_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 18: - mul_mat_vec_iq3_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 19: - mul_mat_vec_iq1_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 20: - mul_mat_vec_iq4_nl_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 21: - mul_mat_vec_iq3_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 22: - mul_mat_vec_iq2_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 23: - mul_mat_vec_iq4_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - case 29: - mul_mat_vec_iq1_m_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, vecs, stream); - break; - } - }); - return Y; -} - -torch::stable::Tensor ggml_mul_mat_a8(torch::stable::Tensor W, // quant weight - torch::stable::Tensor X, // input - int64_t type, int64_t row) { - int col = X.sizes()[1]; - int padded = (col + 512 - 1) / 512 * 512; - int batch = X.sizes()[0]; - const torch::stable::accelerator::DeviceGuard device_guard( - X.get_device_index()); - auto Y = torch::stable::empty({batch, row}, X.scalar_type(), std::nullopt, - W.device()); - cudaStream_t stream = get_current_cuda_stream(); - auto quant_X = torch::stable::empty({batch, padded / 32 * 9}, - torch::headeronly::ScalarType::Int, - std::nullopt, W.device()); - VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), - col, batch, stream); - - switch (type) { - case 2: - ggml_mul_mat_q4_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 3: - ggml_mul_mat_q4_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 6: - ggml_mul_mat_q5_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 7: - ggml_mul_mat_q5_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 8: - ggml_mul_mat_q8_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 10: - ggml_mul_mat_q2_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 11: - ggml_mul_mat_q3_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 12: - ggml_mul_mat_q4_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 13: - ggml_mul_mat_q5_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - case 14: - ggml_mul_mat_q6_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), col, row, batch, padded, row, stream); - break; - } - }); - return Y; -} - -torch::stable::Tensor ggml_moe_a8(torch::stable::Tensor X, // input - torch::stable::Tensor W, // expert weights - torch::stable::Tensor sorted_token_ids, - torch::stable::Tensor expert_ids, - torch::stable::Tensor num_tokens_post_padded, - int64_t type, int64_t row, int64_t top_k, - int64_t tokens) { - int col = X.sizes()[1]; - int padded = (col + 512 - 1) / 512 * 512; - const torch::stable::accelerator::DeviceGuard device_guard( - X.get_device_index()); - auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(), - std::nullopt, W.device()); - cudaStream_t stream = get_current_cuda_stream(); - auto quant_X = torch::stable::empty({tokens, padded / 32 * 9}, - torch::headeronly::ScalarType::Int, - std::nullopt, W.device()); - VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), - col, tokens, stream); - switch (type) { - case 2: - ggml_moe_q4_0_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 3: - ggml_moe_q4_1_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 6: - ggml_moe_q5_0_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 7: - ggml_moe_q5_1_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 8: - ggml_moe_q8_0_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 10: - ggml_moe_q2_K_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 11: - ggml_moe_q3_K_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 12: - ggml_moe_q4_K_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 13: - ggml_moe_q5_K_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - case 14: - ggml_moe_q6_K_q8_1_cuda( - (void*)quant_X.data_ptr(), (void*)W.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(), - (int*)expert_ids.data_ptr(), - (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, - tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream); - break; - } - }); - return Y; -} - -torch::stable::Tensor ggml_moe_a8_vec( - torch::stable::Tensor X, // input - torch::stable::Tensor W, // expert weights - torch::stable::Tensor topk_ids, int64_t top_k, int64_t type, int64_t row, - int64_t tokens) { - int col = X.sizes()[1]; - const int padded = (col + 512 - 1) / 512 * 512; - const torch::stable::accelerator::DeviceGuard device_guard( - X.get_device_index()); - auto Y = torch::stable::empty({tokens * top_k, row}, X.scalar_type(), - std::nullopt, W.device()); - torch::stable::fill_(Y, 0.0); - cudaStream_t stream = get_current_cuda_stream(); - auto quant_X = torch::stable::empty({tokens, padded / 32 * 9}, - torch::headeronly::ScalarType::Int, - std::nullopt, W.device()); - VLLM_STABLE_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { - quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), - (void*)quant_X.data_ptr(), col, tokens, - stream); - switch (type) { - case 2: - moe_vec_q4_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 3: - moe_vec_q4_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 6: - moe_vec_q5_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 7: - moe_vec_q5_1_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 8: - moe_vec_q8_0_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 10: - moe_vec_q2_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 11: - moe_vec_q3_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 12: - moe_vec_q4_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 13: - moe_vec_q5_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 14: - moe_vec_q6_K_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 16: - moe_vec_iq2_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 17: - moe_vec_iq2_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 18: - moe_vec_iq3_xxs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 19: - moe_vec_iq1_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 20: - moe_vec_iq4_nl_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 21: - moe_vec_iq3_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 22: - moe_vec_iq2_s_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 23: - moe_vec_iq4_xs_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - case 29: - moe_vec_iq1_m_q8_1_cuda( - (void*)W.data_ptr(), (void*)quant_X.data_ptr(), - (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens, - col, row, quant_X.stride(0), stream); - break; - } - }); - return Y; -} - -int64_t ggml_moe_get_block_size(int64_t type) { - switch (type) { - case 2: - return MOE_X_Q4_0; - case 3: - return MOE_X_Q4_1; - case 6: - return MOE_X_Q5_0; - case 7: - return MOE_X_Q5_1; - case 8: - return MOE_X_Q8_0; - case 10: - return MOE_X_Q2_K; - case 11: - return MOE_X_Q3_K; - case 12: - return MOE_X_Q4_K; - case 13: - return MOE_X_Q5_K; - case 14: - return MOE_X_Q6_K; - } - return 0; -} diff --git a/csrc/libtorch_stable/quantization/gguf/mmq.cuh b/csrc/libtorch_stable/quantization/gguf/mmq.cuh deleted file mode 100644 index 7c89918c23d..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/mmq.cuh +++ /dev/null @@ -1,610 +0,0 @@ -// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu -template -static __device__ __forceinline__ void mul_mat_q( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - - const block_q_t * x = (const block_q_t *) vx; - const block_q8_1 * y = (const block_q8_1 *) vy; - - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_col_y = nrows_y / QK8_1; - const int blocks_per_warp = WARP_SIZE_GGUF / qi; - - const int & ncols_dst = ncols_y; - - const auto row_dst_0 = blockIdx.x*mmq_y; - const int & row_x_0 = row_dst_0; - - const auto col_dst_0 = blockIdx.y*mmq_x; - const int & col_y_0 = col_dst_0; - - int * tile_x_ql = nullptr; - half2 * tile_x_dm = nullptr; - int * tile_x_qh = nullptr; - int * tile_x_sc = nullptr; - - allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); - - __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; - __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF/QI8_1]; - - float sum[mmq_y/WARP_SIZE_GGUF][mmq_x/nwarps] = {{0.0f}}; - - for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { - - load_tiles(x + row_x_0*blocks_per_row_x + ib0, tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, - threadIdx.y, nrows_x-row_x_0-1, threadIdx.x, blocks_per_row_x); - -#pragma unroll - for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp/qr < blocks_per_row_x; ++ir) { - const auto kqs = ir*WARP_SIZE_GGUF + threadIdx.x; - const int kbxd = kqs / QI8_1; - -#pragma unroll - for (int i = 0; i < mmq_x; i += nwarps) { - const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y-1); // to prevent out-of-bounds memory accesses - const block_q8_1 * by0 = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + kbxd]; - const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; - tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); - } - -#pragma unroll - for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) { - const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF/QI8_1)) % mmq_x; - const auto kby = threadIdx.x % (WARP_SIZE_GGUF/QI8_1); - const int col_y_eff = min(col_y_0 + ids, ncols_y-1); - - // if the sum is not needed it's faster to transform the scale to f32 ahead of time - const half2 * dsi_src = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + ir*(WARP_SIZE_GGUF/QI8_1) + kby].ds; - half2 * dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF/QI8_1) + kby]; - if (need_sum) { - *dsi_dst = *dsi_src; - } else { - float * dfi_dst = (float *) dsi_dst; - *dfi_dst = __low2float(*dsi_src); - } - } - - __syncthreads(); - -// #pragma unroll // unrolling this loop causes too much register pressure - for (int k = ir*WARP_SIZE_GGUF/qr; k < (ir+1)*WARP_SIZE_GGUF/qr; k += vdr) { -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - sum[i/WARP_SIZE_GGUF][j/nwarps] += vec_dot( - tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, - threadIdx.x + i, threadIdx.y + j, k); - } - } - } - __syncthreads(); - } - } - -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { - const auto col_dst = col_dst_0 + j + threadIdx.y; - if (col_dst >= ncols_dst) { - return; - } - -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - const auto row_dst = row_dst_0 + threadIdx.x + i; - if (row_dst >= nrows_dst) { - continue; - } - dst[col_dst*nrows_dst + row_dst] = sum[i/WARP_SIZE_GGUF][j/nwarps]; - } - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_0 64 -#define MMQ_Y_Q4_0 128 -#define NWARPS_Q4_0 8 -#else -#define MMQ_X_Q4_0 4 -#define MMQ_Y_Q4_0 32 -#define NWARPS_Q4_0 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_0, 2) -#endif -mul_mat_q4_0( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_0; - const int mmq_y = MMQ_Y_Q4_0; - const int nwarps = NWARPS_Q4_0; - - mul_mat_q, - load_tiles_q4_0, VDR_Q4_0_Q8_1_MMQ, vec_dot_q4_0_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_0_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - - int mmq_x = MMQ_X_Q4_0; - int mmq_y = MMQ_Y_Q4_0; - int nwarps = NWARPS_Q4_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_1 64 -#define MMQ_Y_Q4_1 128 -#define NWARPS_Q4_1 8 -#else -#define MMQ_X_Q4_1 4 -#define MMQ_Y_Q4_1 32 -#define NWARPS_Q4_1 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_1, 2) -#endif -mul_mat_q4_1( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_1; - const int mmq_y = MMQ_Y_Q4_1; - const int nwarps = NWARPS_Q4_1; - - mul_mat_q, - load_tiles_q4_1, VDR_Q4_1_Q8_1_MMQ, vec_dot_q4_1_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_1_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - - int mmq_x = MMQ_X_Q4_1; - int mmq_y = MMQ_Y_Q4_1; - int nwarps = NWARPS_Q4_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_1<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_1<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_0 64 -#define MMQ_Y_Q5_0 128 -#define NWARPS_Q5_0 8 -#else -#define MMQ_X_Q5_0 4 -#define MMQ_Y_Q5_0 32 -#define NWARPS_Q5_0 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_0, 2) -#endif -mul_mat_q5_0( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_0; - const int mmq_y = MMQ_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - mul_mat_q, - load_tiles_q5_0, VDR_Q5_0_Q8_1_MMQ, vec_dot_q5_0_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_0_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - - const int mmq_x = MMQ_X_Q5_0; - const int mmq_y = MMQ_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_1 64 -#define MMQ_Y_Q5_1 128 -#define NWARPS_Q5_1 8 -#else -#define MMQ_X_Q5_1 4 -#define MMQ_Y_Q5_1 32 -#define NWARPS_Q5_1 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_1, 2) -#endif -mul_mat_q5_1( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_1; - const int mmq_y = MMQ_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - mul_mat_q, - load_tiles_q5_1, VDR_Q5_1_Q8_1_MMQ, vec_dot_q5_1_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_1_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - const int mmq_x = MMQ_X_Q5_1; - const int mmq_y = MMQ_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_1<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_1<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q8_0 64 -#define MMQ_Y_Q8_0 128 -#define NWARPS_Q8_0 8 -#else -#define MMQ_X_Q8_0 4 -#define MMQ_Y_Q8_0 32 -#define NWARPS_Q8_0 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q8_0, 2) -#endif -mul_mat_q8_0( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q8_0; - const int mmq_y = MMQ_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - mul_mat_q, - load_tiles_q8_0, VDR_Q8_0_Q8_1_MMQ, vec_dot_q8_0_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q8_0_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - const int mmq_x = MMQ_X_Q8_0; - const int mmq_y = MMQ_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q8_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q8_0<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q2_K 64 -#define MMQ_Y_Q2_K 128 -#define NWARPS_Q2_K 8 -#else -#define MMQ_X_Q2_K 4 -#define MMQ_Y_Q2_K 32 -#define NWARPS_Q2_K 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q2_K, 2) -#endif -mul_mat_q2_K( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q2_K; - const int mmq_y = MMQ_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - mul_mat_q, - load_tiles_q2_K, VDR_Q2_K_Q8_1_MMQ, vec_dot_q2_K_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q2_K_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - const int mmq_x = MMQ_X_Q2_K; - const int mmq_y = MMQ_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q2_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q2_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q3_K 64 -#define MMQ_Y_Q3_K 128 -#define NWARPS_Q3_K 8 -#else -#define MMQ_X_Q3_K 4 -#define MMQ_Y_Q3_K 32 -#define NWARPS_Q3_K 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q3_K, 2) -#endif -mul_mat_q3_K( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - - const int mmq_x = MMQ_X_Q3_K; - const int mmq_y = MMQ_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - mul_mat_q, - load_tiles_q3_K, VDR_Q3_K_Q8_1_MMQ, vec_dot_q3_K_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q3_K_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - - const int mmq_x = MMQ_X_Q3_K; - const int mmq_y = MMQ_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q3_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q3_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q4_K 64 -#define MMQ_Y_Q4_K 128 -#define NWARPS_Q4_K 8 -#else -#define MMQ_X_Q4_K 4 -#define MMQ_Y_Q4_K 32 -#define NWARPS_Q4_K 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q4_K, 2) -#endif -mul_mat_q4_K( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q4_K; - const int mmq_y = MMQ_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - mul_mat_q, - load_tiles_q4_K, VDR_Q4_K_Q8_1_MMQ, vec_dot_q4_K_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q4_K_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - const int mmq_x = MMQ_X_Q4_K; - const int mmq_y = MMQ_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q4_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q4_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q5_K 64 -#define MMQ_Y_Q5_K 128 -#define NWARPS_Q5_K 8 -#else -#define MMQ_X_Q5_K 4 -#define MMQ_Y_Q5_K 32 -#define NWARPS_Q5_K 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q5_K, 2) -#endif -mul_mat_q5_K( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q5_K; - const int mmq_y = MMQ_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - mul_mat_q, - load_tiles_q5_K, VDR_Q5_K_Q8_1_MMQ, vec_dot_q5_K_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q5_K_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - - const int mmq_x = MMQ_X_Q5_K; - const int mmq_y = MMQ_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q5_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q5_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} - -#if defined(USE_ROCM) -#define MMQ_X_Q6_K 64 -#define MMQ_Y_Q6_K 128 -#define NWARPS_Q6_K 8 -#else -#define MMQ_X_Q6_K 4 -#define MMQ_Y_Q6_K 32 -#define NWARPS_Q6_K 4 -#endif - -template static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF*NWARPS_Q6_K, 2) -#endif -mul_mat_q6_K( - const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { - const int mmq_x = MMQ_X_Q6_K; - const int mmq_y = MMQ_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - mul_mat_q, - load_tiles_q6_K, VDR_Q6_K_Q8_1_MMQ, vec_dot_q6_K_q8_1_mul_mat> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); -} - -template -static void ggml_mul_mat_q6_K_q8_1_cuda( - const void * vx, const void * vy, scalar_t * dst, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, cudaStream_t stream) { - const int mmq_x = MMQ_X_Q6_K; - const int mmq_y = MMQ_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - const bool need_check = false; - mul_mat_q6_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } else { - const bool need_check = true; - mul_mat_q6_K<<>> - (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); - } -} diff --git a/csrc/libtorch_stable/quantization/gguf/mmvq.cuh b/csrc/libtorch_stable/quantization/gguf/mmvq.cuh deleted file mode 100644 index e27bec7af5b..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/mmvq.cuh +++ /dev/null @@ -1,212 +0,0 @@ -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu -template -static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, const int ncols, const int nrows, const int nvecs) { - const auto row = blockIdx.x*blockDim.y + threadIdx.y; - const auto vec = blockIdx.y; - - if (row >= nrows || vec >= nvecs) { - return; - } - - const int blocks_per_row = ncols / qk; - const int blocks_per_warp = vdr * WARP_SIZE / qi; - const int nrows_y = (ncols + 512 - 1) / 512 * 512; - - - // partial sum for each thread - float tmp = 0.0f; - - const block_q_t * x = (const block_q_t *) vx; - const block_q8_1 * y = (const block_q8_1 *) vy; - - for (auto i = threadIdx.x / (qi/vdr); i < blocks_per_row; i += blocks_per_warp) { - const int ibx = row*blocks_per_row + i; // x block index - - const int iby = vec*(nrows_y/QK8_1) + i * (qk/QK8_1); // y block index that aligns with ibx - - const int iqs = vdr * (threadIdx.x % (qi/vdr)); // x block quant index when casting the quants to int - - tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); - } - - // sum up partial sums and write back result -#pragma unroll - for (int mask = WARP_SIZE/2; mask > 0; mask >>= 1) { - tmp += VLLM_SHFL_XOR_SYNC(tmp, mask); - } - - if (threadIdx.x == 0) { - dst[vec*nrows + row] = tmp; - } -} - -template -static void mul_mat_vec_q4_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q4_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q8_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q2_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q3_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q4_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q5_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_q6_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq2_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq3_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq1_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq1_m_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq4_nl_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq4_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} - -template -static void mul_mat_vec_iq3_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, nvecs, 1); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - mul_mat_vec_q - <<>>(vx, vy, dst, ncols, nrows, nvecs); -} diff --git a/csrc/libtorch_stable/quantization/gguf/moe.cuh b/csrc/libtorch_stable/quantization/gguf/moe.cuh deleted file mode 100644 index a2f9f46c8f8..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/moe.cuh +++ /dev/null @@ -1,739 +0,0 @@ -#include - -/* Adapted from ./csrc/quantization/gguf/mmq.cuh - based on ./vllm/model_executor/layers/fused_moe/experts/triton_moe.py */ -template -static __device__ __forceinline__ void moe_q( - const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* __restrict__ sorted_token_ids, - const int* __restrict__ expert_ids, - const int* __restrict__ num_tokens_post_padded, const int exp_stride, - const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, - const int nrows_dst, const int top_k) { - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_col_y = nrows_y / QK8_1; - const int blocks_per_warp = WARP_SIZE_GGUF / qi; - - const int ncols_dst = ncols_y * top_k; - - const auto row_dst_0 = blockIdx.x * mmq_y; - const int& row_x_0 = row_dst_0; - - const auto col_dst_0 = blockIdx.y * mmq_x; - - int token_offs[mmq_x / nwarps]; - for (int i = 0; i < mmq_x; i += nwarps) { - token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i]; - } - - const int exp_idx = expert_ids[blockIdx.y]; - if (exp_idx > 255 || exp_idx < 0) return; - if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return; - - const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride); - const block_q8_1* y = (const block_q8_1*)(vy); - - int* tile_x_ql = nullptr; - half2* tile_x_dm = nullptr; - int* tile_x_qh = nullptr; - int* tile_x_sc = nullptr; - - allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); - - __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; - __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1]; - - float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}}; - - for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { - load_tiles(x + row_x_0 * blocks_per_row_x + ib0, tile_x_ql, tile_x_dm, - tile_x_qh, tile_x_sc, threadIdx.y, nrows_x - row_x_0 - 1, - threadIdx.x, blocks_per_row_x); - - const int n_per_r = ((qk * blocks_per_warp) / qr); -#pragma unroll - for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) { - const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x; - const int kbxd = kqs / QI8_1; - -#pragma unroll - for (int i = 0; i < mmq_x; i += nwarps) { - const int col_y_eff = token_offs[i / nwarps] / top_k; - const int block_x = ib0 * (qk / QK8_1) + kbxd; - if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { - const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + block_x]; - const int index_y = - (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; - tile_y_qs[index_y] = - get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); - } - } - - if (threadIdx.x < n_per_r / QK8_1) { - const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); - const int col_y_eff = token_offs[threadIdx.y] / top_k; - const int block_x = - ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby; - - if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { - const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds; - half2* dsi_dst = - &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby]; - - if (need_sum) { - *dsi_dst = *dsi_src; - } else { - float* dfi_dst = (float*)dsi_dst; - *dfi_dst = __low2float(*dsi_src); - } - } - } - __syncthreads(); - - // #pragma unroll // unrolling this loop causes too much register pressure - for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; - k += vdr) { -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - sum[i / WARP_SIZE_GGUF][j / nwarps] += - vec_dot(tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, - tile_y_ds, threadIdx.x + i, threadIdx.y + j, k); - } - } - } - __syncthreads(); - } - } - -#pragma unroll - for (int j = 0; j < mmq_x; j += nwarps) { - const int col_dst = token_offs[j / nwarps]; - if (col_dst >= ncols_dst) { - return; - } - -#pragma unroll - for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { - const auto row_dst = row_dst_0 + threadIdx.x + i; - if (row_dst >= nrows_dst) { - continue; - } - dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps]; - } - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q4_0 8 - #define MOE_Y_Q4_0 128 - #define NWARPS_Q4_0 8 -#else - #define MOE_X_Q4_0 4 - #define MOE_Y_Q4_0 32 - #define NWARPS_Q4_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) -#endif - moe_q4_0(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_0; - const int mmq_y = MOE_Y_Q4_0; - const int nwarps = NWARPS_Q4_0; - - moe_q, load_tiles_q4_0, - VDR_Q4_0_Q8_1_MMQ, vec_dot_q4_0_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q4_0_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - int mmq_x = MOE_X_Q4_0; - int mmq_y = MOE_Y_Q4_0; - int nwarps = NWARPS_Q4_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q4_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q4_1 8 - #define MOE_Y_Q4_1 128 - #define NWARPS_Q4_1 8 -#else - #define MOE_X_Q4_1 4 - #define MOE_Y_Q4_1 32 - #define NWARPS_Q4_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) -#endif - moe_q4_1(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_1; - const int mmq_y = MOE_Y_Q4_1; - const int nwarps = NWARPS_Q4_1; - - moe_q, load_tiles_q4_1, - VDR_Q4_1_Q8_1_MMQ, vec_dot_q4_1_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q4_1_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - int mmq_x = MOE_X_Q4_1; - int mmq_y = MOE_Y_Q4_1; - int nwarps = NWARPS_Q4_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_1<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q4_1<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q5_0 8 - #define MOE_Y_Q5_0 128 - #define NWARPS_Q5_0 8 -#else - #define MOE_X_Q5_0 4 - #define MOE_Y_Q5_0 32 - #define NWARPS_Q5_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) -#endif - moe_q5_0(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_0; - const int mmq_y = MOE_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - moe_q, load_tiles_q5_0, - VDR_Q5_0_Q8_1_MMQ, vec_dot_q5_0_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q5_0_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_0; - const int mmq_y = MOE_Y_Q5_0; - const int nwarps = NWARPS_Q5_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q5_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q5_1 8 - #define MOE_Y_Q5_1 128 - #define NWARPS_Q5_1 8 -#else - #define MOE_X_Q5_1 4 - #define MOE_Y_Q5_1 32 - #define NWARPS_Q5_1 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) -#endif - moe_q5_1(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_1; - const int mmq_y = MOE_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - moe_q, load_tiles_q5_1, - VDR_Q5_1_Q8_1_MMQ, vec_dot_q5_1_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q5_1_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_1; - const int mmq_y = MOE_Y_Q5_1; - const int nwarps = NWARPS_Q5_1; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_1<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q5_1<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q8_0 8 - #define MOE_Y_Q8_0 128 - #define NWARPS_Q8_0 8 -#else - #define MOE_X_Q8_0 4 - #define MOE_Y_Q8_0 32 - #define NWARPS_Q8_0 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) -#endif - moe_q8_0(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q8_0; - const int mmq_y = MOE_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - moe_q, load_tiles_q8_0, - VDR_Q8_0_Q8_1_MMQ, vec_dot_q8_0_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q8_0_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q8_0; - const int mmq_y = MOE_Y_Q8_0; - const int nwarps = NWARPS_Q8_0; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q8_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q8_0<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q2_K 8 - #define MOE_Y_Q2_K 128 - #define NWARPS_Q2_K 8 -#else - #define MOE_X_Q2_K 4 - #define MOE_Y_Q2_K 32 - #define NWARPS_Q2_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) -#endif - moe_q2_K(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q2_K; - const int mmq_y = MOE_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - moe_q, load_tiles_q2_K, - VDR_Q2_K_Q8_1_MMQ, vec_dot_q2_K_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q2_K_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q2_K; - const int mmq_y = MOE_Y_Q2_K; - const int nwarps = NWARPS_Q2_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q2_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q2_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q3_K 8 - #define MOE_Y_Q3_K 128 - #define NWARPS_Q3_K 8 -#else - #define MOE_X_Q3_K 4 - #define MOE_Y_Q3_K 32 - #define NWARPS_Q3_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) -#endif - moe_q3_K(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - - const int mmq_x = MOE_X_Q3_K; - const int mmq_y = MOE_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - moe_q, load_tiles_q3_K, - VDR_Q3_K_Q8_1_MMQ, vec_dot_q3_K_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} -template -static void ggml_moe_q3_K_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q3_K; - const int mmq_y = MOE_Y_Q3_K; - const int nwarps = NWARPS_Q3_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q3_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q3_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q4_K 8 - #define MOE_Y_Q4_K 128 - #define NWARPS_Q4_K 8 -#else - #define MOE_X_Q4_K 4 - #define MOE_Y_Q4_K 32 - #define NWARPS_Q4_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) -#endif - moe_q4_K(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q4_K; - const int mmq_y = MOE_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - moe_q, load_tiles_q4_K, - VDR_Q4_K_Q8_1_MMQ, vec_dot_q4_K_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q4_K_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q4_K; - const int mmq_y = MOE_Y_Q4_K; - const int nwarps = NWARPS_Q4_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q4_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q4_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q5_K 8 - #define MOE_Y_Q5_K 128 - #define NWARPS_Q5_K 8 -#else - #define MOE_X_Q5_K 4 - #define MOE_Y_Q5_K 32 - #define NWARPS_Q5_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) -#endif - moe_q5_K(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q5_K; - const int mmq_y = MOE_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - moe_q, load_tiles_q5_K, - VDR_Q5_K_Q8_1_MMQ, vec_dot_q5_K_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q5_K_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q5_K; - const int mmq_y = MOE_Y_Q5_K; - const int nwarps = NWARPS_Q5_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q5_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q5_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} - -#if defined(USE_ROCM) - #define MOE_X_Q6_K 8 - #define MOE_Y_Q6_K 128 - #define NWARPS_Q6_K 8 -#else - #define MOE_X_Q6_K 4 - #define MOE_Y_Q6_K 32 - #define NWARPS_Q6_K 4 -#endif - -template -static __global__ void -#if defined(USE_ROCM) -__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) -#endif - moe_q6_K(const void* __restrict__ vx, const void* __restrict__ vy, - scalar_t* __restrict__ dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, - const int top_k) { - const int mmq_x = MOE_X_Q6_K; - const int mmq_y = MOE_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - moe_q, load_tiles_q6_K, - VDR_Q6_K_Q8_1_MMQ, vec_dot_q6_K_q8_1_mul_mat>( - vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); -} - -template -static void ggml_moe_q6_K_q8_1_cuda( - const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids, - const int* expert_ids, const int* num_tokens_post_padded, - const int exp_stride, const int ncols_x, const int nrows_x, - const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k, - const int tokens_post_padded, cudaStream_t stream) { - const int mmq_x = MOE_X_Q6_K; - const int mmq_y = MOE_Y_Q6_K; - const int nwarps = NWARPS_Q6_K; - - const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; - const int block_num_y = (tokens_post_padded) / mmq_x; - const dim3 block_nums(block_num_x, block_num_y, 1); - const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); - - if (nrows_x % mmq_y == 0) { - constexpr bool need_check = false; - moe_q6_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } else { - constexpr bool need_check = true; - moe_q6_K<<>>( - w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded, - exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k); - } -} diff --git a/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh b/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh deleted file mode 100644 index 60f65a1bfdc..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/moe_vec.cuh +++ /dev/null @@ -1,338 +0,0 @@ -// copied and adapted from -// https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu -template -static __global__ void moe_vec_q(const void* __restrict__ vx, - const void* __restrict__ vy, - scalar_t* __restrict__ dst, - const int* topk_ids, const int topk, - const int ncols, const int nrows, - const int token_stride) { - const auto row = blockIdx.x * blockDim.y + threadIdx.y; - - const auto token = blockIdx.z / topk; - const auto expert = (topk_ids)[blockIdx.z]; - - if (row >= nrows) { - return; - } - - const int blocks_per_row = ncols / qk; - const int blocks_per_warp = vdr * WARP_SIZE / qi; - - // partial sum for each thread - float tmp = 0.0f; - - const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row; - const block_q8_1* y = - (const block_q8_1*)(((const int*)vy) + token * token_stride); - - for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; - i += blocks_per_warp) { - const int ibx = row * blocks_per_row + i; // x block index - - const int iby = i * (qk / QK8_1); // y block index that aligns with ibx - - const int iqs = - vdr * - (threadIdx.x % - (qi / vdr)); // x block quant index when casting the quants to int - - tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); - } - - // sum up partial sums and write back result -#pragma unroll - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - tmp += VLLM_SHFL_XOR_SYNC(tmp, mask); - } - - if (threadIdx.x == 0) { - dst[blockIdx.z * nrows + row] = tmp; - } -} - -template -static void moe_vec_q4_0_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_1_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_0_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_1_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q8_0_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q2_K_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q3_K_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_K_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_K_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q6_K_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_xxs_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_xs_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_s_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq3_xxs_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_s_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_m_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq4_nl_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q<<>>( - vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq4_xs_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} - -template -static void moe_vec_iq3_s_q8_1_cuda(const void* vx, const void* vy, - scalar_t* dst, const int* topk_ids, - const int top_k, const int tokens, - const int ncols, const int nrows, - const int token_stride, - cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, - ncols, nrows, token_stride); -} diff --git a/csrc/libtorch_stable/quantization/gguf/vecdotq.cuh b/csrc/libtorch_stable/quantization/gguf/vecdotq.cuh deleted file mode 100644 index d0d4c74ed37..00000000000 --- a/csrc/libtorch_stable/quantization/gguf/vecdotq.cuh +++ /dev/null @@ -1,1812 +0,0 @@ -// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/vecdotq.cuh -// and https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu -static __device__ __forceinline__ int get_int_b2(const void * x, const int & i32) { - const uint16_t * x16 = (const uint16_t *) x; // assume at least 2 byte alignment - - int x32 = x16[2*i32 + 0] << 0; - x32 |= x16[2*i32 + 1] << 16; - - return x32; -} - -static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32) { - return ((const int *) x)[i32]; // assume at least 4 byte alignment -} - -static __device__ __forceinline__ int get_int_from_int8(const int8_t * x8, const int & i32) { - const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; -} - -static __device__ __forceinline__ int get_int_from_uint8(const uint8_t * x8, const int & i32) { - const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment - int x32 = 0; - x32 |= x16[0] << 0; - x32 |= x16[1] << 16; - return x32; -} - -static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t * x8, const int & i32) { - return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment -} - -static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t * x8, const int & i32) { - return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment -} - -// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called -// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q - -#define VDR_Q4_0_Q8_1_MMVQ 2 -#define VDR_Q4_0_Q8_1_MMQ 4 - -template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( - const int * v, const int * u, const float & d4, const half2 & ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; - const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; - - // SIMD dot product of quantized values - sumi = __dp4a(vi0, u[2*i+0], sumi); - sumi = __dp4a(vi1, u[2*i+1], sumi); - } - - const float2 ds8f = __half22float2(ds8); - - // second part effectively subtracts 8 from each quant value - return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y); -#endif -} - -#define VDR_Q4_1_Q8_1_MMVQ 2 -#define VDR_Q4_1_Q8_1_MMQ 4 - -template static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl( - const int * v, const int * u, const half2 & dm4, const half2 & ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; - const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; - - // SIMD dot product of quantized values - sumi = __dp4a(vi0, u[2*i+0], sumi); - sumi = __dp4a(vi1, u[2*i+1], sumi); - } - - const float2 tmp = __half22float2(__hmul2(dm4, ds8)); - const float d4d8 = tmp.x; - const float m4s8 = tmp.y; - - // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it - return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); -#endif -} - -#define VDR_Q5_0_Q8_1_MMVQ 2 -#define VDR_Q5_0_Q8_1_MMQ 4 - -template static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl( - const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits - vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 - vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 - vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 - vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 - sumi = __dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values - - int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits - vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 - vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 - vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 - vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 - sumi = __dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values - } - - const float2 ds8f = __half22float2(ds8); - - // second part effectively subtracts 16 from each quant value - return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y); -#endif -} - - -#define VDR_Q5_1_Q8_1_MMVQ 2 -#define VDR_Q5_1_Q8_1_MMQ 4 - -template static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl( - const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits - vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 - vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 - vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 - vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 - sumi = __dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values - - int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits - vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 - vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 - vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 - vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 - sumi = __dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values - } - - const float2 tmp = __half22float2(__hmul2(dm5, ds8)); - const float d5d8 = tmp.x; - const float m5s8 = tmp.y; - - // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it - return sumi*d5d8 + m5s8 / (QI5_1 / vdr); -#endif -} - -#define VDR_Q8_0_Q8_1_MMVQ 2 -#define VDR_Q8_0_Q8_1_MMQ 8 - -template static __device__ __forceinline__ float vec_dot_q8_0_q8_1_impl( - const int * v, const int * u, const float & d8_0, const float & d8_1) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - // SIMD dot product of quantized values - sumi = __dp4a(v[i], u[i], sumi); - } - return d8_0*d8_1 * sumi; -#endif -} - -template static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl( - const int * v, const int * u, const half2 & dm8, const half2 & ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - int sumi = 0; - -#pragma unroll - for (int i = 0; i < vdr; ++i) { - // SIMD dot product of quantized values - sumi = __dp4a(v[i], u[i], sumi); - } - - const float2 tmp = __half22float2(__hmul2(dm8, ds8)); - const float d8d8 = tmp.x; - const float m8s8 = tmp.y; - - // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it - return sumi*d8d8 + m8s8 / (QI8_1 / vdr); -#endif -} - -#define VDR_Q2_K_Q8_1_MMVQ 1 -#define VDR_Q2_K_Q8_1_MMQ 2 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( - const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, - const half2 & dm2, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR2_K; ++i) { - const int sc = scales[2*i]; - - const int vi = (v >> (2*i)) & 0x03030303; - - sumf_d += d8[i] * (__dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product - - // fill int with 4x m - int m = sc >> 4; - m |= m << 8; - m |= m << 16; - sumf_m += d8[i] * __dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values - } - - const float2 dm2f = __half22float2(dm2); - - return dm2f.x*sumf_d - dm2f.y*sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( - const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ scales, - const half2 & dm2, const float & d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi_d = 0; - int sumi_m = 0; - -#pragma unroll - for (int i0 = 0; i0 < QI8_1; i0 += QI8_1/2) { - int sumi_d_sc = 0; - - const int sc = scales[i0 / (QI8_1/2)]; - - // fill int with 4x m - int m = sc >> 4; - m |= m << 8; - m |= m << 16; - -#pragma unroll - for (int i = i0; i < i0 + QI8_1/2; ++i) { - sumi_d_sc = __dp4a(v[i], u[i], sumi_d_sc); // SIMD dot product - sumi_m = __dp4a(m, u[i], sumi_m); // multiply sum of q8_1 values with m - } - - sumi_d += sumi_d_sc * (sc & 0xF); - } - - const float2 dm2f = __half22float2(dm2); - - return d8 * (dm2f.x*sumi_d - dm2f.y*sumi_m); -#endif -} - -#define VDR_Q3_K_Q8_1_MMVQ 1 -#define VDR_Q3_K_Q8_1_MMQ 2 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( - const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, - const int & scale_offset, const float & d3, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - float sumf = 0.0f; - -#pragma unroll - for (int i = 0; i < QR3_K; ++i) { - const int isc = scale_offset + 2*i; - - const int isc_low = isc % (QK_K/32); - const int sc_shift_low = 4 * (isc / (QK_K/32)); - const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; - - const int isc_high = isc % (QK_K/64); - const int sc_shift_high = 2 * (isc / (QK_K/64)); - const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; - - const int sc = (sc_low | sc_high) - 32; - - const int vil = (vl >> (2*i)) & 0x03030303; - - const int vih = ((vh >> i) << 2) & 0x04040404; - - const int vi = __vsubss4(vil, vih); - - sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product - } - - return d3 * sumf; -#endif -} - -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( - const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales, - const float & d3, const float & d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - int sumi = 0; - -#pragma unroll - for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) { - int sumi_sc = 0; - - for (int i = i0; i < i0 + QI8_1/2; ++i) { - sumi_sc = __dp4a(v[i], u[i], sumi_sc); // SIMD dot product - } - - sumi += sumi_sc * scales[i0 / (QI8_1/2)]; - } - - return d3*d8 * sumi; -#endif -} - -#define VDR_Q4_K_Q8_1_MMVQ 2 -#define VDR_Q4_K_Q8_1_MMQ 8 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( - const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, - const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR4_K; ++i) { - const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F; - const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F; - - const int dot1 = __dp4a(v1i, u[2*i+1], __dp4a(v0i, u[2*i+0], 0)); // SIMD dot product - const int dot2 = __dp4a(0x01010101, u[2*i+1], __dp4a(0x01010101, u[2*i+0], 0)); // sum of u - - sumf_d += d8[i] * (dot1 * sc[i]); - sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values - } - - const float2 dm4f = __half22float2(dm4); - return dm4f.x*sumf_d - dm4f.y*sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( - const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, - const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) { - int sumi_d = 0; - -#pragma unroll - for (int j = 0; j < QI8_1; ++j) { - sumi_d = __dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product - } - - const float2 ds8f = __half22float2(ds8[i]); - - sumf_d += ds8f.x * (sc[i] * sumi_d); - sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val - } - - const float2 dm4f = __half22float2(dm4); - - return dm4f.x*sumf_d - dm4f.y*sumf_m; -#endif -} - -#define VDR_Q5_K_Q8_1_MMVQ 2 -#define VDR_Q5_K_Q8_1_MMQ 8 - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( - const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc, - const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR5_K; ++i) { - const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F; - const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F; - - const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; - const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; - - const int v0i = vl0i | vh0i; - const int v1i = vl1i | vh1i; - - const int dot1 = __dp4a(v0i, u[2*i+0], __dp4a(v1i, u[2*i+1], 0)); // SIMD dot product - const int dot2 = __dp4a(0x01010101, u[2*i+0], __dp4a(0x01010101, u[2*i+1], 0)); // sum of u - - sumf_d += d8[i] * (dot1 * sc[i]); - sumf_m += d8[i] * (dot2 * m[i]); - } - - const float2 dm5f = __half22float2(dm5); - return dm5f.x*sumf_d - dm5f.y*sumf_m; -#endif -} - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( - const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, - const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - float sumf_d = 0.0f; - float sumf_m = 0.0f; - -#pragma unroll - for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) { - int sumi_d = 0; - -#pragma unroll - for (int j = 0; j < QI8_1; ++j) { - sumi_d = __dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product - } - - const float2 ds8f = __half22float2(ds8[i]); - - sumf_d += ds8f.x * (sc[i] * sumi_d); - sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val - } - - const float2 dm4f = __half22float2(dm4); - - return dm4f.x*sumf_d - dm4f.y*sumf_m; -#endif -} - -#define VDR_Q6_K_Q8_1_MMVQ 1 -#define VDR_Q6_K_Q8_1_MMQ 8 - -// contiguous v/x values -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( - const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, - const float & d, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - float sumf = 0.0f; - -#pragma unroll - for (int i = 0; i < QR6_K; ++i) { - const int sc = scales[4*i]; - const int vil = (vl >> (4*i)) & 0x0F0F0F0F; - const int vih = ((vh >> (4*i)) << 4) & 0x30303030; - const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 - - sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product - } - - return d*sumf; -#endif -} - -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( - const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc, - const float & d6, const float * __restrict__ d8) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - float sumf_d = 0.0f; - -#pragma unroll - for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { - int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale - -#pragma unroll - for (int i = i0; i < i0 + 2; ++i) { - sumi_d.x = __dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product - sumi_d.x = __dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product - - sumi_d.y = __dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product - sumi_d.y = __dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product - } - - sumf_d += d8[i0/4] * (sc[i0/2+0]*sumi_d.x + sc[i0/2+1]*sumi_d.y); - } - - return d6 * sumf_d; -#endif -} - -static __device__ __forceinline__ float vec_dot_q4_0_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq; - - int v[VDR_Q4_0_Q8_1_MMVQ]; - int u[2*VDR_Q4_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); - u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); - } - - return vec_dot_q4_0_q8_1_impl(v, u, __half2float(bq4_0->d), bq8_1->ds); -} - -template static __device__ __forceinline__ void allocate_tiles_q4_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF/QI4_0) + mmq_y/QI4_0]; - *x_ql = tile_x_qs; - *x_dm = (half2 *) tile_x_d; -} - -template static __device__ __forceinline__ void load_tiles_q4_0( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI4_0; - const int kqsx = k % QI4_0; - - const block_q4_0 * bx0 = (const block_q4_0 *) vx; - float * x_dmf = (float *) x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q4_0 * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); - // x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i / QI4_0 + kbx] = bxi->d; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_0; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_0) { - int i = i0 + i_offset * QI4_0 + k / blocks_per_tile_x_row; - if (need_check) { - i = min(i, i_max); - } - const block_q4_0 * bxi = bx0 + i*blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i / QI4_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q4_0_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - (void)x_qh; (void)x_sc; - - const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); - const float * x_dmf = (const float *) x_dm; - - int u[2*VDR_Q4_0_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) { - u[2*l+0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2*l+1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_0) % WARP_SIZE_GGUF]; - } - - return vec_dot_q4_0_q8_1_impl - (&x_ql[i * (WARP_SIZE_GGUF + 1) + k], u, x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i/QI4_0 + k/QI4_0], - y_ds[j * (WARP_SIZE_GGUF/QI8_1) + (2*k/QI8_1) % (WARP_SIZE_GGUF/QI8_1)]); -} - -static __device__ __forceinline__ float vec_dot_q4_1_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq; - - int v[VDR_Q4_1_Q8_1_MMVQ]; - int u[2*VDR_Q4_1_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); - u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); - } - - return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); -} - -template static __device__ __forceinline__ void allocate_tiles_q4_1(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI4_1) + mmq_y/QI4_1]; - *x_ql = tile_x_qs; - *x_dm = tile_x_dm; -} - -template static __device__ __forceinline__ void load_tiles_q4_1( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI4_1; - const int kqsx = k % QI4_1; - - const block_q4_1 * bx0 = (const block_q4_1 *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q4_1 * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_1; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_1) { - int i = i0 + i_offset * QI4_1 + k / blocks_per_tile_x_row; - if (need_check) { - i = min(i, i_max); - } - const block_q4_1 * bxi = bx0 + i*blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF/QI4_1) + i / QI4_1 + kbxd] = bxi->dm; - } -} - -static __device__ __forceinline__ float vec_dot_q4_1_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); - - int u[2*VDR_Q4_1_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) { - u[2*l+0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2*l+1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_1) % WARP_SIZE_GGUF]; - } - - return vec_dot_q4_1_q8_1_impl - (&x_ql[i * (WARP_SIZE_GGUF + 1) + k], u, x_dm[i * (WARP_SIZE_GGUF/QI4_1) + i/QI4_1 + k/QI4_1], - y_ds[j * (WARP_SIZE_GGUF/QI8_1) + (2*k/QI8_1) % (WARP_SIZE_GGUF/QI8_1)]); -} - -static __device__ __forceinline__ float vec_dot_q5_0_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq; - - int vl[VDR_Q5_0_Q8_1_MMVQ]; - int vh[VDR_Q5_0_Q8_1_MMVQ]; - int u[2*VDR_Q5_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { - vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); - vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); - u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); - } - - return vec_dot_q5_0_q8_1_impl(vl, vh, u, __half2float(bq5_0->d), bq8_1->ds); -} - -template static __device__ __forceinline__ void allocate_tiles_q5_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF/QI5_0) + mmq_y/QI5_0]; - - *x_ql = tile_x_ql; - *x_dm = (half2 *) tile_x_d; -} - -template static __device__ __forceinline__ void load_tiles_q5_0( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI5_0; - const int kqsx = k % QI5_0; - - const block_q5_0 * bx0 = (const block_q5_0 *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q5_0 * bxi = bx0 + i*blocks_per_row + kbx; - const int ql = get_int_from_uint8(bxi->qs, kqsx); - const int qh = get_int_from_uint8(bxi->qh, 0) >> (4 * (k % QI5_0)); - - int qs0 = (ql >> 0) & 0x0F0F0F0F; - qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 - qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 - qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 - qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 - qs0 = __vsubss4(qs0, 0x10101010); // subtract 16 - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2*k+0] = qs0; - - int qs1 = (ql >> 4) & 0x0F0F0F0F; - qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 - qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 - qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 - qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 - qs1 = __vsubss4(qs1, 0x10101010); // subtract 16 - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2*k+1] = qs1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_0; - const int kbxd = k % blocks_per_tile_x_row; - float * x_dmf = (float *) x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_0) { - int i = i0 + i_offset * QI5_0 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_0 * bxi = bx0 + i*blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF/QI5_0) + i / QI5_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q5_0_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); - const int index_bx = i * (WARP_SIZE_GGUF/QI5_0) + i/QI5_0 + k/QI5_0; - const float * x_dmf = (const float *) x_dm; - const float * y_df = (const float *) y_ds; - - int u[2*VDR_Q5_0_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q5_0_Q8_1_MMQ; ++l) { - u[2*l+0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2*l+1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_0) % WARP_SIZE_GGUF]; - } - - return vec_dot_q8_0_q8_1_impl - (&x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2 * k], u, x_dmf[index_bx], y_df[j * (WARP_SIZE_GGUF/QI8_1) + (2*k/QI8_1) % (WARP_SIZE_GGUF/QI8_1)]); -} - -static __device__ __forceinline__ float vec_dot_q5_1_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq; - - int vl[VDR_Q5_1_Q8_1_MMVQ]; - int vh[VDR_Q5_1_Q8_1_MMVQ]; - int u[2*VDR_Q5_1_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { - vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); - vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); - u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); - } - - return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); -} - -template static __device__ __forceinline__ void allocate_tiles_q5_1(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI5_1) + mmq_y/QI5_1]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; -} - -template static __device__ __forceinline__ void load_tiles_q5_1( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI5_1; - const int kqsx = k % QI5_1; - - const block_q5_1 * bx0 = (const block_q5_1 *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_1 * bxi = bx0 + i*blocks_per_row + kbx; - - const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); - const int qh = get_int_from_uint8_aligned(bxi->qh, 0) >> (4 * (k % QI5_1)); - - int qs0 = (ql >> 0) & 0x0F0F0F0F; - qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 - qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 - qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 - qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2*k+0] = qs0; - - int qs1 = (ql >> 4) & 0x0F0F0F0F; - qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 - qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 - qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 - qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2*k+1] = qs1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_1; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_1) { - int i = i0 + i_offset * QI5_1 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_1 * bxi = bx0 + i*blocks_per_row + kbxd; - - x_dm[i * (WARP_SIZE_GGUF/QI5_1) + i / QI5_1 + kbxd] = bxi->dm; - } -} - -static __device__ __forceinline__ float vec_dot_q5_1_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); - const int index_bx = i * (WARP_SIZE_GGUF/QI5_1) + + i/QI5_1 + k/QI5_1; - - int u[2*VDR_Q5_1_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < VDR_Q5_1_Q8_1_MMQ; ++l) { - u[2*l+0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; - u[2*l+1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_1) % WARP_SIZE_GGUF]; - } - - return vec_dot_q8_1_q8_1_impl - (&x_ql[i * (2*WARP_SIZE_GGUF + 1) + 2 * k], u, x_dm[index_bx], y_ds[j * (WARP_SIZE_GGUF/QI8_1) + (2*k/QI8_1) % (WARP_SIZE_GGUF/QI8_1)]); -} - -static __device__ __forceinline__ float vec_dot_q8_0_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq; - - int v[VDR_Q8_0_Q8_1_MMVQ]; - int u[VDR_Q8_0_Q8_1_MMVQ]; - -#pragma unroll - for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { - v[i] = get_int_from_int8(bq8_0->qs, iqs + i); - u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); - } - - return vec_dot_q8_0_q8_1_impl(v, u, __half2float(bq8_0->d), __low2float(bq8_1->ds)); -} - -template static __device__ __forceinline__ void allocate_tiles_q8_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF/QI8_0) + mmq_y/QI8_0]; - - *x_ql = tile_x_qs; - *x_dm = (half2 *) tile_x_d; -} - -template static __device__ __forceinline__ void load_tiles_q8_0( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI8_0; - const int kqsx = k % QI8_0; - float * x_dmf = (float *) x_dm; - - const block_q8_0 * bx0 = (const block_q8_0 *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q8_0 * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_int8(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI8_0; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI8_0) { - int i = i0 + i_offset * QI8_0 + k / blocks_per_tile_x_row; - - if (need_check) { - i = min(i, i_max); - } - const block_q8_0 * bxi = bx0 + i*blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF/QI8_0) + i / QI8_0 + kbxd] = __half2float(bxi->d); - } -} - -static __device__ __forceinline__ float vec_dot_q8_0_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const float * x_dmf = (const float *) x_dm; - const float * y_df = (const float *) y_ds; - - return vec_dot_q8_0_q8_1_impl - (&x_ql[i * (WARP_SIZE_GGUF + 1) + k], &y_qs[j * WARP_SIZE_GGUF + k], x_dmf[i * (WARP_SIZE_GGUF/QI8_0) + i/QI8_0 + k/QI8_0], - y_df[j * (WARP_SIZE_GGUF/QI8_1) + k/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_q2_K_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q2_K * bq2_K = (const block_q2_K *) vbq; - - const int bq8_offset = QR2_K * (iqs / QI8_1); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); - - const uint8_t * scales = bq2_K->scales + scale_offset; - - const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); - int u[QR2_K]; - float d8[QR2_K]; - -#pragma unroll - for (int i = 0; i < QR2_K; ++ i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + i].ds); - } - - return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); -} - -template static __device__ __forceinline__ void allocate_tiles_q2_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI2_K) + mmq_y/QI2_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF/4) + mmq_y/4]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template static __device__ __forceinline__ void load_tiles_q2_K( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI2_K; - const int kqsx = k % QI2_K; - - const block_q2_K * bx0 = (const block_q2_K *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI2_K; - const int kbxd = k % blocks_per_tile_x_row; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI2_K) { - int i = (i0 + i_offset * QI2_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K * bxi = bx0 + i*blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF/QI2_K) + i / QI2_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { - int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF/4); - - if (need_check) { - i = min(i, i_max); - } - const block_q2_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/4)) / (QI2_K/4); - x_sc[i * (WARP_SIZE_GGUF/4) + i / 4 + k % (WARP_SIZE_GGUF/4)] = get_int_from_uint8_aligned(bxi->scales, k % (QI2_K/4)); - } -} - -static __device__ __forceinline__ float vec_dot_q2_K_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const int kbx = k / QI2_K; - const int ky = (k % QI2_K) * QR2_K; - const float * y_df = (const float *) y_ds; - - int v[QR2_K*VDR_Q2_K_Q8_1_MMQ]; - - const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx*QI2_K + (QI2_K/2) * (ky/(2*QI2_K)) + ky % (QI2_K/2); - const int shift = 2 * ((ky % (2*QI2_K)) / (QI2_K/2)); - -#pragma unroll - for (int l = 0; l < QR2_K*VDR_Q2_K_Q8_1_MMQ; ++l) { - v[l] = (x_ql[kqsx + l] >> shift) & 0x03030303; - } - - const uint8_t * scales = ((const uint8_t *) &x_sc[i * (WARP_SIZE_GGUF/4) + i/4 + kbx*4]) + ky/4; - - const int index_y = j * WARP_SIZE_GGUF + (QR2_K*k) % WARP_SIZE_GGUF; - return vec_dot_q2_K_q8_1_impl_mmq(v, &y_qs[index_y], scales, x_dm[i * (WARP_SIZE_GGUF/QI2_K) + i/QI2_K + kbx], y_df[index_y/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_q3_K_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q3_K * bq3_K = (const block_q3_K *) vbq; - - const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); - - const float d = __half2float(bq3_K->d); - - const int vl = get_int_from_uint8(bq3_K->qs, iqs); - - // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted - const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; - - int u[QR3_K]; - float d8[QR3_K]; - -#pragma unroll - for (int i = 0; i < QR3_K; ++i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + i].ds); - } - - return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); -} - -template static __device__ __forceinline__ void allocate_tiles_q3_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI3_K) + mmq_y/QI3_K]; - __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF/2) + mmq_y/2]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF/4) + mmq_y/4]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_qh = tile_x_qh; - *x_sc = tile_x_sc; -} - -template static __device__ __forceinline__ void load_tiles_q3_K( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI3_K; - const int kqsx = k % QI3_K; - - const block_q3_K * bx0 = (const block_q3_K *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - if (need_check) { - i = min(i, i_max); - } - const block_q3_K * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI3_K; - const int kbxd = k % blocks_per_tile_x_row; - float * x_dmf = (float *) x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI3_K) { - int i = (i0 + i_offset * QI3_K + k / blocks_per_tile_x_row) % mmq_y; - if (need_check) { - i = min(i, i_max); - } - const block_q3_K * bxi = bx0 + i*blocks_per_row + kbxd; - x_dmf[i * (WARP_SIZE_GGUF/QI3_K) + i / QI3_K + kbxd] = __half2float(bxi->d); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 2) { - int i = i0 + i_offset * 2 + k / (WARP_SIZE_GGUF/2); - if (need_check) { - i = min(i, i_max); - } - const block_q3_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/2)) / (QI3_K/2); - // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted - x_qh[i * (WARP_SIZE_GGUF/2) + i / 2 + k % (WARP_SIZE_GGUF/2)] = ~get_int_from_uint8(bxi->hmask, k % (QI3_K/2)); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { - int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF/4); - if (need_check) { - i = min(i, i_max); - } - const block_q3_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/4)) / (QI3_K/4); - - const int ksc = k % (QI3_K/4); - - const int ksc_low = ksc % (QI3_K/8); - const int shift_low = 4 * (ksc / (QI3_K/8)); - const int sc_low = (get_int_from_uint8(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; - - const int ksc_high = QI3_K/8; - const int shift_high = 2 * ksc; - const int sc_high = ((get_int_from_uint8(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; - - const int sc = __vsubss4(sc_low | sc_high, 0x20202020); - - x_sc[i * (WARP_SIZE_GGUF/4) + i / 4 + k % (WARP_SIZE_GGUF/4)] = sc; - } -} - -static __device__ __forceinline__ float vec_dot_q3_K_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - - const int kbx = k / QI3_K; - const int ky = (k % QI3_K) * QR3_K; - const float * x_dmf = (const float *) x_dm; - const float * y_df = (const float *) y_ds; - - const int8_t * scales = ((const int8_t *) (x_sc + i * (WARP_SIZE_GGUF/4) + i/4 + kbx*4)) + ky/4; - - int v[QR3_K*VDR_Q3_K_Q8_1_MMQ]; - -#pragma unroll - for (int l = 0; l < QR3_K*VDR_Q3_K_Q8_1_MMQ; ++l) { - const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx*QI3_K + (QI3_K/2) * (ky/(2*QI3_K)) + ky % (QI3_K/2); - const int shift = 2 * ((ky % 32) / 8); - const int vll = (x_ql[kqsx + l] >> shift) & 0x03030303; - - const int vh = x_qh[i * (WARP_SIZE_GGUF/2) + i/2 + kbx * (QI3_K/2) + (ky+l)%8] >> ((ky+l) / 8); - const int vlh = (vh << 2) & 0x04040404; - - v[l] = __vsubss4(vll, vlh); - } - - const int index_y = j * WARP_SIZE_GGUF + (k*QR3_K) % WARP_SIZE_GGUF; - return vec_dot_q3_K_q8_1_impl_mmq(v, &y_qs[index_y], scales, x_dmf[i * (WARP_SIZE_GGUF/QI3_K) + i/QI3_K + kbx], y_df[index_y/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_q4_K_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - const block_q4_K * bq4_K = (const block_q4_K *) vbq; - - int v[2]; - int u[2*QR4_K]; - float d8[QR4_K]; - - // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 - const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2)); - - // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 - // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 - // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 - // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 - - const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); - v[0] = q4[0]; - v[1] = q4[4]; - - const uint16_t * scales = (const uint16_t *)bq4_K->scales; - uint16_t aux[2]; - const int j = bq8_offset/2; - if (j < 2) { - aux[0] = scales[j+0] & 0x3f3f; - aux[1] = scales[j+2] & 0x3f3f; - } else { - aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); - aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); - } - const uint8_t * sc = (const uint8_t *)aux; - const uint8_t * m = sc + 2; - - for (int i = 0; i < QR4_K; ++i) { - const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; - d8[i] = __low2float(bq8i->ds); - - const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); - u[2*i+0] = q8[0]; - u[2*i+1] = q8[4]; - } - - return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); -} - -template static __device__ __forceinline__ void allocate_tiles_q4_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI4_K) + mmq_y/QI4_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF/8) + mmq_y/8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template static __device__ __forceinline__ void load_tiles_q4_K( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI4_K; // == 0 if QK_K == 256 - const int kqsx = k % QI4_K; // == k if QK_K == 256 - - const block_q4_K * bx0 = (const block_q4_K *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - const block_q4_K * bxi = bx0 + i*blocks_per_row + kbx; - x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_K) { - int i = (i0 + i_offset * QI4_K + k / blocks_per_tile_x_row) % mmq_y; - if (need_check) { - i = min(i, i_max); - } - const block_q4_K * bxi = bx0 + i*blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF/QI4_K) + i / QI4_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF/8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q4_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/8)) / (QI4_K/8); - - const int * scales = (const int *) bxi->scales; - - const int ksc = k % (WARP_SIZE_GGUF/8); - // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 - int scales8 = (scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F; // lower 4 bits - scales8 |= (scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits - - x_sc[i * (WARP_SIZE_GGUF/8) + i / 8 + ksc] = scales8; - } -} - -static __device__ __forceinline__ float vec_dot_q4_K_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - (void)x_qh; - - const uint8_t * sc = ((const uint8_t *) &x_sc[i * (WARP_SIZE_GGUF/8) + i/8 + k/16]) + 2*((k % 16) / 8); - - const int index_y = j * WARP_SIZE_GGUF + (QR4_K*k) % WARP_SIZE_GGUF; - return vec_dot_q4_K_q8_1_impl_mmq(&x_ql[i * (WARP_SIZE_GGUF + 1) + k], &y_qs[index_y], sc, sc+8, - x_dm[i * (WARP_SIZE_GGUF/QI4_K) + i/QI4_K], &y_ds[index_y/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q5_K * bq5_K = (const block_q5_K *) vbq; - - int vl[2]; - int vh[2]; - int u[2*QR5_K]; - float d8[QR5_K]; - - const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2)); - const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); - const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4)); - - vl[0] = ql[0]; - vl[1] = ql[4]; - - vh[0] = qh[0] >> bq8_offset; - vh[1] = qh[4] >> bq8_offset; - - const uint16_t * scales = (const uint16_t *)bq5_K->scales; - uint16_t aux[2]; - const int j = bq8_offset/2; - if (j < 2) { - aux[0] = scales[j+0] & 0x3f3f; - aux[1] = scales[j+2] & 0x3f3f; - } else { - aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); - aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); - } - const uint8_t * sc = (const uint8_t *)aux; - const uint8_t * m = sc + 2; - -#pragma unroll - for (int i = 0; i < QR5_K; ++i) { - const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; - d8[i] = __low2float(bq8i->ds); - - const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); - u[2*i+0] = q8[0]; - u[2*i+1] = q8[4]; - } - - return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); -} - -template static __device__ __forceinline__ void allocate_tiles_q5_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI5_K) + mmq_y/QI5_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF/8) + mmq_y/8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template static __device__ __forceinline__ void load_tiles_q5_K( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI5_K; // == 0 if QK_K == 256 - const int kqsx = k % QI5_K; // == k if QK_K == 256 - - const block_q5_K * bx0 = (const block_q5_K *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K * bxi = bx0 + i*blocks_per_row + kbx; - const int ky = QR5_K*kqsx; - - const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); - const int ql0 = (ql >> 0) & 0x0F0F0F0F; - const int ql1 = (ql >> 4) & 0x0F0F0F0F; - - const int qh = get_int_from_uint8_aligned(bxi->qh, kqsx % (QI5_K/4)); - const int qh0 = ((qh >> (2 * (kqsx / (QI5_K/4)) + 0)) << 4) & 0x10101010; - const int qh1 = ((qh >> (2 * (kqsx / (QI5_K/4)) + 1)) << 4) & 0x10101010; - - const int kq0 = ky - ky % (QI5_K/2) + k % (QI5_K/4) + 0; - const int kq1 = ky - ky % (QI5_K/2) + k % (QI5_K/4) + (QI5_K/4); - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + kq0] = ql0 | qh0; - x_ql[i * (2*WARP_SIZE_GGUF + 1) + kq1] = ql1 | qh1; - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_K) { - int i = (i0 + i_offset * QI5_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K * bxi = bx0 + i*blocks_per_row + kbxd; - x_dm[i * (WARP_SIZE_GGUF/QI5_K) + i / QI5_K + kbxd] = bxi->dm; - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF/8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q5_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/8)) / (QI5_K/8); - - const int * scales = (const int *) bxi->scales; - - const int ksc = k % (WARP_SIZE_GGUF/8); - - // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 - int scales8 = (scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F; // lower 4 bits - scales8 |= (scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits - - x_sc[i * (WARP_SIZE_GGUF/8) + i / 8 + ksc] = scales8; - } -} - -static __device__ __forceinline__ float vec_dot_q5_K_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const uint8_t * sc = ((const uint8_t *) &x_sc[i * (WARP_SIZE_GGUF/8) + i/8 + k/16]) + 2 * ((k % 16) / 8); - - const int index_x = i * (QR5_K*WARP_SIZE_GGUF + 1) + QR5_K*k; - const int index_y = j * WARP_SIZE_GGUF + (QR5_K*k) % WARP_SIZE_GGUF; - return vec_dot_q5_K_q8_1_impl_mmq(&x_ql[index_x], &y_qs[index_y], sc, sc+8, - x_dm[i * (WARP_SIZE_GGUF/QI5_K) + i/QI5_K], &y_ds[index_y/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_q6_K_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - - const block_q6_K * bq6_K = (const block_q6_K *) vbq; - - const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4); - const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); - const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); - - const int vl = get_int_from_uint8(bq6_K->ql, iqs); - const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; - - const int8_t * scales = bq6_K->scales + scale_offset; - - int u[QR6_K]; - float d8[QR6_K]; - -#pragma unroll - for (int i = 0; i < QR6_K; ++i) { - u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); - d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); - } - - return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, __half2float(bq6_K->d), d8); -} - -template static __device__ __forceinline__ void allocate_tiles_q6_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { - __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE_GGUF) + mmq_y]; - __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF/QI6_K) + mmq_y/QI6_K]; - __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF/8) + mmq_y/8]; - - *x_ql = tile_x_ql; - *x_dm = tile_x_dm; - *x_sc = tile_x_sc; -} - -template static __device__ __forceinline__ void load_tiles_q6_K( - const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, - int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { - const int kbx = k / QI6_K; // == 0 if QK_K == 256 - const int kqsx = k % QI6_K; // == k if QK_K == 256 - - const block_q6_K * bx0 = (const block_q6_K *) vx; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { - int i = i0 + i_offset; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K * bxi = bx0 + i*blocks_per_row + kbx; - const int ky = QR6_K*kqsx; - - const int ql = get_int_from_uint8(bxi->ql, kqsx); - const int ql0 = (ql >> 0) & 0x0F0F0F0F; - const int ql1 = (ql >> 4) & 0x0F0F0F0F; - - const int qh = get_int_from_uint8(bxi->qh, (QI6_K/4) * (kqsx / (QI6_K/2)) + kqsx % (QI6_K/4)); - const int qh0 = ((qh >> (2 * ((kqsx % (QI6_K/2)) / (QI6_K/4)))) << 4) & 0x30303030; - const int qh1 = (qh >> (2 * ((kqsx % (QI6_K/2)) / (QI6_K/4)))) & 0x30303030; - - const int kq0 = ky - ky % QI6_K + k % (QI6_K/2) + 0; - const int kq1 = ky - ky % QI6_K + k % (QI6_K/2) + (QI6_K/2); - - x_ql[i * (2*WARP_SIZE_GGUF + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020); - x_ql[i * (2*WARP_SIZE_GGUF + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020); - } - - const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI6_K; // == 1 if QK_K == 256 - const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 - float * x_dmf = (float *) x_dm; - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI6_K) { - int i = (i0 + i_offset * QI6_K + k / blocks_per_tile_x_row) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K * bxi = bx0 + i*blocks_per_row + kbxd; - - x_dmf[i * (WARP_SIZE_GGUF/QI6_K) + i / QI6_K + kbxd] = __half2float(bxi->d); - } - -#pragma unroll - for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { - int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF/8)) % mmq_y; - - if (need_check) { - i = min(i, i_max); - } - - const block_q6_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE_GGUF/8)) / 4; - - x_sc[i * (WARP_SIZE_GGUF/8) + i / 8 + k % (WARP_SIZE_GGUF/8)] = get_int_from_int8(bxi->scales, k % (QI6_K/8)); - } -} - -static __device__ __forceinline__ float vec_dot_q6_K_q8_1_mul_mat( - const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, - const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { - const float * x_dmf = (const float *) x_dm; - const float * y_df = (const float *) y_ds; - - const int8_t * sc = ((const int8_t *) &x_sc[i * (WARP_SIZE_GGUF/8) + i/8 + k/8]); - - const int index_x = i * (QR6_K*WARP_SIZE_GGUF + 1) + QR6_K*k; - const int index_y = j * WARP_SIZE_GGUF + (QR6_K*k) % WARP_SIZE_GGUF; - return vec_dot_q6_K_q8_1_impl_mmq(&x_ql[index_x], &y_qs[index_y], sc, x_dmf[i * (WARP_SIZE_GGUF/QI6_K) + i/QI6_K], &y_df[index_y/QI8_1]); -} - -static __device__ __forceinline__ float vec_dot_iq2_xxs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - const block_iq2_xxs * bq2 = (const block_iq2_xxs *) vbq; - - const int ib32 = iqs; - const uint16_t * q2 = bq2->qs + 4*ib32; - const uint8_t * aux8 = (const uint8_t *)q2; - const int8_t * q8 = bq8_1[ib32].qs; - uint32_t aux32 = q2[2] | (q2[3] << 16); - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); - const uint8_t signs = ksigns_iq2xs[aux32 & 127]; - for (int j = 0; j < 8; ++j) { - sumi += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - aux32 >>= 7; - } - const float d = __half2float(bq2->d) * (0.5f + aux32) * __half2float(bq8_1[ib32].ds.x) * 0.25f; - return d * sumi; -} - -static __device__ __forceinline__ float vec_dot_iq2_xs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { - const block_iq2_xs * bq2 = (const block_iq2_xs *) vbq; - - const int ib32 = iqs; - const uint16_t * q2 = bq2->qs + 4*ib32; - const int8_t * q8 = bq8_1[ib32].qs; - const uint8_t ls1 = bq2->scales[ib32] & 0xf; - const uint8_t ls2 = bq2->scales[ib32] >> 4; - int sumi1 = 0; - for (int l = 0; l < 2; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi1 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - int sumi2 = 0; - for (int l = 2; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi2 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - const float d = __half2float(bq2->d) * __half2float(bq8_1[ib32].ds.x) * 0.25f; - return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); -} - -static __device__ __forceinline__ float vec_dot_iq2_s_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - const block_iq2_s * bq2 = (const block_iq2_s *) vbq; - - const int ib32 = iqs; - const int8_t * q8 = bq8_1[ib32].qs; - const uint8_t * signs = bq2->qs + QK_K/8 + 4*ib32; - const uint8_t ls1 = bq2->scales[ib32] & 0xf; - const uint8_t ls2 = bq2->scales[ib32] >> 4; - int sumi1 = 0; - for (int l = 0; l < 2; ++l) { - const uint32_t * grid = (const uint32_t *)(iq2s_grid + (bq2->qs[4*ib32+l] | ((bq2->qh[ib32] << (8-2*l)) & 0x300))); - const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid[1] ^ signs1, signs1); - sumi1 = __dp4a(grid_l, *((const int *)q8 + 0), sumi1); - sumi1 = __dp4a(grid_h, *((const int *)q8 + 1), sumi1); - q8 += 8; - } - int sumi2 = 0; - for (int l = 2; l < 4; ++l) { - const uint32_t * grid = (const uint32_t *)(iq2s_grid + (bq2->qs[4*ib32+l] | ((bq2->qh[ib32] << (8-2*l)) & 0x300))); - const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid[1] ^ signs1, signs1); - sumi2 = __dp4a(grid_l, *((const int *)q8 + 0), sumi2); - sumi2 = __dp4a(grid_h, *((const int *)q8 + 1), sumi2); - q8 += 8; - } - const float d = __half2float(bq2->d) * __low2float(bq8_1[ib32].ds) * 0.25f; - return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); -#endif -} - -static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - const block_iq3_xxs * bq2 = (const block_iq3_xxs *) vbq; - - const int ib32 = iqs; - const uint8_t * q3 = bq2->qs + 8*ib32; - const uint16_t * gas = (const uint16_t *)(bq2->qs + QK_K/4) + 2*ib32; - const int8_t * q8 = bq8_1[ib32].qs; - uint32_t aux32 = gas[0] | (gas[1] << 16); - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint32_t * grid1 = iq3xxs_grid + q3[2*l+0]; - const uint32_t * grid2 = iq3xxs_grid + q3[2*l+1]; - const uint32_t * signs = (const uint32_t *)(ksigns64 + (aux32 & 127)); - const int grid_l = __vsub4(grid1[0] ^ signs[0], signs[0]); - const int grid_h = __vsub4(grid2[0] ^ signs[1], signs[1]); - sumi = __dp4a(grid_l, *((int *)q8+0), sumi); - sumi = __dp4a(grid_h, *((int *)q8+1), sumi); - q8 += 8; - aux32 >>= 7; - } - const float d = __half2float(bq2->d) * (0.5f + aux32) * __low2float(bq8_1[ib32].ds) * 0.5f; - return d * sumi; -#endif -} - -static __device__ __forceinline__ float vec_dot_iq3_s_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - const block_iq3_s * bq2 = (const block_iq3_s *) vbq; - - const int ib32 = iqs; - const uint8_t * qs = bq2->qs + 8*ib32; - const int8_t * q8 = bq8_1[ib32].qs; - int sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint32_t * grid1 = iq3xs_grid + (qs[2*l+0] | ((bq2->qh[ib32] << (8 - 2*l)) & 256)); - const uint32_t * grid2 = iq3xs_grid + (qs[2*l+1] | ((bq2->qh[ib32] << (7 - 2*l)) & 256)); - uint32_t signs0 = __vcmpeq4(((bq2->signs[4*ib32+l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); - uint32_t signs1 = __vcmpeq4(((bq2->signs[4*ib32+l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); - const int grid_l = __vsub4(grid1[0] ^ signs0, signs0); - const int grid_h = __vsub4(grid2[0] ^ signs1, signs1); - sumi = __dp4a(grid_l, *((int *)q8+0), sumi); - sumi = __dp4a(grid_h, *((int *)q8+1), sumi); - q8 += 8; - } - const float d = __half2float(bq2->d) * (0.5f + ((bq2->scales[ib32/2] >> 4*(ib32%2)) & 0xf)) * __low2float(bq8_1[ib32].ds) * 0.5f; - return d * sumi; -#endif -} - -static __device__ __forceinline__ float vec_dot_iq1_s_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - const block_iq1_s * bq1 = (const block_iq1_s *) vbq; - - const int qs_packed = get_int_b2(bq1->qs, iqs); - const uint8_t * qs = (const uint8_t *) &qs_packed; - - const int qh = bq1->qh[iqs]; - - int sumi = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int grid = iq1s_grid_gpu[qs[l0/2] | (((qh >> 3*(l0/2)) & 0x07) << 8)]; - - const int grid0 = (grid >> 0) & 0x0F0F0F0F; - const int grid1 = (grid >> 4) & 0x0F0F0F0F; - - const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); - - sumi = __dp4a(grid0, u0, sumi); - sumi = __dp4a(grid1, u1, sumi); - } - - const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); - const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000); - const float2 ds = __half22float2(bq8_1[iqs].ds); - return d1q * (ds.x*sumi + ds.y*delta); -#endif -} - -static __device__ __forceinline__ float vec_dot_iq1_m_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - const block_iq1_m * bq1 = (const block_iq1_m *) vbq; - - const int qs_packed = get_int_b4(bq1->qs, iqs); - const uint8_t * qs = (const uint8_t *) &qs_packed; - - int sumi[2] = {0}; - float sumf[2] = {0.0f}; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int qhl = bq1->qh[2*iqs + l0/4] >> (4 * ((l0/2) % 2)); - - const int grid = iq1s_grid_gpu[qs[l0/2] | ((qhl & 0x07) << 8)]; - - const int grid0 = (grid >> 0) & 0x0F0F0F0F; - const int grid1 = (grid >> 4) & 0x0F0F0F0F; - - const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); - const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); - - sumi[l0/4] = __dp4a(grid0, u0, sumi[l0/4]); - sumi[l0/4] = __dp4a(grid1, u1, sumi[l0/4]); - - const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f*IQ1M_DELTA/0x08); - int sumy = 0; - sumy = __dp4a(u0, 0x01010101, sumy); - sumy = __dp4a(u1, 0x01010101, sumy); - sumf[l0/4] += delta*sumy; - } - - const uint16_t * sc = (const uint16_t *) bq1->scales; - - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); - const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); - - const int tmp = sc[iqs/2] >> (6*(iqs%2)); - const int sc0 = 2*((tmp >> 0) & 0x07) + 1; - const int sc1 = 2*((tmp >> 3) & 0x07) + 1; - return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); -#endif -} - -static __device__ __forceinline__ void get_int_from_table_16(const uint32_t & q4, const uint8_t * values, - int & val1, int & val2) { - - uint32_t aux32; const uint8_t * q8 = (const uint8_t *)&aux32; - aux32 = q4 & 0x0f0f0f0f; - uint16_t v1 = values[q8[0]] | (values[q8[1]] << 8); - uint16_t v2 = values[q8[2]] | (values[q8[3]] << 8); - val1 = v1 | (v2 << 16); - aux32 = (q4 >> 4) & 0x0f0f0f0f; - v1 = values[q8[0]] | (values[q8[1]] << 8); - v2 = values[q8[2]] | (values[q8[3]] << 8); - val2 = v1 | (v2 << 16); -} - -static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - - const block_iq4_nl * bq = (const block_iq4_nl *) vbq; - - const uint16_t * q4 = (const uint16_t *)bq->qs + 2*iqs; - const int32_t * q8 = (const int32_t *)bq8_1->qs + iqs; - - const uint8_t * values = (const uint8_t *)kvalues_iq4nl; - - int v1, v2; - int sumi1 = 0, sumi2 = 0; - for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { - const uint32_t aux = q4[2*l] | (q4[2*l+1] << 16); - get_int_from_table_16(aux, values, v1, v2); - sumi1 = __dp4a(v1, q8[l+0], sumi1); - sumi2 = __dp4a(v2, q8[l+4], sumi2); - } - const float d = __half2float(bq->d) * __low2float(bq8_1->ds); - return d * (sumi1 + sumi2); -#endif -} - - -static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { -#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM - const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq; - const uint8_t * values = (const uint8_t *)kvalues_iq4nl; - - // iqs is 0...7 - const int ib32 = iqs; - const int32_t * q8 = (const int *)bq8_1[ib32].qs; - const uint32_t * q4 = (const uint32_t *)bq4->qs + 4*ib32; - const int8_t ls = ((bq4->scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((bq4->scales_h >> 2*ib32) & 3) << 4); - const float d = __half2float(bq4->d) * (ls - 32) * __low2float(bq8_1[ib32].ds); - int v1, v2; - int sumi1 = 0, sumi2 = 0; - for (int j = 0; j < 4; ++j) { - get_int_from_table_16(q4[j], values, v1, v2); - sumi1 = __dp4a(v1, q8[j+0], sumi1); - sumi2 = __dp4a(v2, q8[j+4], sumi2); - } - return d * (sumi1 + sumi2); -#endif -} \ No newline at end of file diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu index b325d30a041..9a00f66241b 100644 --- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu +++ b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu @@ -150,6 +150,8 @@ void rearrange_kn_weight_as_n32k16_order( void* b_zero_reorder = has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr; + const torch::stable::accelerator::DeviceGuard device_guard( + b_qweight.get_device_index()); cudaStream_t stream = get_current_cuda_stream(); if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) { allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>( diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh index ce96c2d11fe..ac33d5f2ce6 100644 --- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh +++ b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh @@ -6,7 +6,7 @@ #include -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" using marlin::MarlinScalarType2; namespace allspark { diff --git a/csrc/quantization/machete/Readme.md b/csrc/libtorch_stable/quantization/machete/Readme.md similarity index 100% rename from csrc/quantization/machete/Readme.md rename to csrc/libtorch_stable/quantization/machete/Readme.md diff --git a/csrc/quantization/machete/generate.py b/csrc/libtorch_stable/quantization/machete/generate.py similarity index 95% rename from csrc/quantization/machete/generate.py rename to csrc/libtorch_stable/quantization/machete/generate.py index e12601e9e97..11a5bbdd13c 100644 --- a/csrc/quantization/machete/generate.py +++ b/csrc/libtorch_stable/quantization/machete/generate.py @@ -39,10 +39,10 @@ namespace machete { {% for impl_config in impl_configs %} {% set type_sig = gen_type_sig(impl_config.types) -%} {% for s in impl_config.schedules %} -extern torch::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs); +extern torch::stable::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs); {%- endfor %} -torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) { +torch::stable::Tensor mm_dispatch_{{type_sig}}(MMArgs args) { [[maybe_unused]] auto M = args.A.size(0); [[maybe_unused]] auto N = args.B.size(1); [[maybe_unused]] auto K = args.A.size(1); @@ -59,14 +59,14 @@ torch::Tensor mm_dispatch_{{type_sig}}(MMArgs args) { if (*args.maybe_schedule == "{{ gen_sch_sig(s) }}") return impl_{{type_sig}}_sch_{{ gen_sch_sig(s) }}(args); {%- endfor %} - TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for " + STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for " "schedule = ", *args.maybe_schedule); } {%- endfor %} -static inline std::optional maybe_scalartype( - std::optional const& t) { +static inline std::optional maybe_scalartype( + std::optional const& t) { if (!t) { return std::nullopt; } else { @@ -74,7 +74,7 @@ static inline std::optional maybe_scalartype( }; } -torch::Tensor mm_dispatch(MMArgs args) { +torch::stable::Tensor mm_dispatch(MMArgs args) { auto out_type = args.maybe_out_type.value_or(args.A.scalar_type()); auto a_type = args.A.scalar_type(); auto maybe_g_scales_type = maybe_scalartype(args.maybe_group_scales); @@ -105,19 +105,19 @@ torch::Tensor mm_dispatch(MMArgs args) { } {%- endfor %} - TORCH_CHECK_NOT_IMPLEMENTED( + STD_TORCH_CHECK_NOT_IMPLEMENTED( false, "machete_mm(..) is not implemented for " - "a_type=", args.A.scalar_type(), + "a_type=", torch::headeronly::toString(args.A.scalar_type()), ", b_type=", args.b_type.str(), - ", out_type=", out_type, + ", out_type=", torch::headeronly::toString(out_type), ", with_group_scale_type=", maybe_g_scales_type - ? toString(*maybe_g_scales_type) : "None", + ? torch::headeronly::toString(*maybe_g_scales_type) : "None", ", with_group_zeropoint_type=", maybe_g_zeros_type - ? toString(*maybe_g_zeros_type) : "None", + ? torch::headeronly::toString(*maybe_g_zeros_type) : "None", ", with_channel_scale_type=", maybe_ch_scales_type - ? toString(*maybe_ch_scales_type) : "None", + ? torch::headeronly::toString(*maybe_ch_scales_type) : "None", ", with_token_scale_type=", maybe_tok_scales_type - ? toString(*maybe_tok_scales_type) : "None", + ? torch::headeronly::toString(*maybe_tok_scales_type) : "None", "; implemented types are: \\n", {%- for impl_config in impl_configs %} {% set t = impl_config.types -%} @@ -197,7 +197,7 @@ using Kernel_{{type_sig}} = MacheteKernelTemplate< {% for sch in schs %} {% set sch_sig = gen_sch_sig(sch) -%} -torch::Tensor +torch::stable::Tensor impl_{{type_sig}}_sch_{{sch_sig}}(MMArgs args) { return run_impl>(args); } @@ -212,7 +212,7 @@ PREPACK_TEMPLATE = """ namespace machete { -torch::Tensor prepack_B_dispatch(PrepackBArgs args) { +torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args) { auto convert_type = args.maybe_group_scales_type.value_or(args.a_type); {%- for t in types %} {% set b_type = unsigned_type_with_bitwidth(t.b_num_bits) %} @@ -231,12 +231,12 @@ torch::Tensor prepack_B_dispatch(PrepackBArgs args) { } {%- endfor %} - TORCH_CHECK_NOT_IMPLEMENTED(false, + STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "prepack_B_dispatch(..) is not implemented for " - "atype = ", args.a_type, + "atype = ", torch::headeronly::toString(args.a_type), ", b_type = ", args.b_type.str(), ", with_group_scales_type= ", args.maybe_group_scales_type ? - toString(*args.maybe_group_scales_type) : "None"); + torch::headeronly::toString(*args.maybe_group_scales_type) : "None"); } }; // namespace machete diff --git a/csrc/quantization/machete/machete_collective_builder.cuh b/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh similarity index 94% rename from csrc/quantization/machete/machete_collective_builder.cuh rename to csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh index ee825583dee..ba8da0af2c3 100644 --- a/csrc/quantization/machete/machete_collective_builder.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh @@ -1,6 +1,6 @@ #pragma once -#include "cutlass_extensions/vllm_collective_builder.cuh" +#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh" #include "machete_mainloop.cuh" namespace cutlass::gemm::collective { diff --git a/csrc/quantization/machete/machete_interleaving_utils.cuh b/csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh similarity index 100% rename from csrc/quantization/machete/machete_interleaving_utils.cuh rename to csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh diff --git a/csrc/quantization/machete/machete_mainloop.cuh b/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh similarity index 100% rename from csrc/quantization/machete/machete_mainloop.cuh rename to csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh diff --git a/csrc/quantization/machete/machete_mm_kernel.cuh b/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh similarity index 86% rename from csrc/quantization/machete/machete_mm_kernel.cuh rename to csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh index cc50e68b058..57655f369cd 100644 --- a/csrc/quantization/machete/machete_mm_kernel.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh @@ -1,8 +1,6 @@ #pragma once -#include -#include -#include +#include // clang-format off // The cutlass include order matters (annoyingly) @@ -20,9 +18,9 @@ // clang-format on #include "cutlass_extensions/cute_utils.cuh" -#include "cutlass_extensions/vllm_numeric_conversion.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "cutlass_extensions/torch_utils.hpp" +#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" #include "machete_collective_builder.cuh" #include "machete_prepacked_layout.cuh" #include "machete_interleaving_utils.cuh" @@ -175,19 +173,23 @@ struct MacheteKernelTemplate { static Arguments create_arguments( cudaStream_t stream, - torch::Tensor const& A, // MxK matrix - torch::Tensor const& B, // KxN prepacked matrix - torch::Tensor& D, // MxN matrix - std::optional const& maybe_g_scales, // scale_KxN matrix - std::optional const& maybe_g_zeros, // scale_KxN matrix + torch::stable::Tensor const& A, // MxK matrix + torch::stable::Tensor const& B, // KxN prepacked matrix + torch::stable::Tensor& D, // MxN matrix + std::optional const& + maybe_g_scales, // scale_KxN matrix + std::optional const& + maybe_g_zeros, // scale_KxN matrix std::optional maybe_group_size, - std::optional const& maybe_ch_scales, // len N vector - std::optional const& maybe_tok_scales) // len M vector + std::optional const& + maybe_ch_scales, // len N vector + std::optional const& + maybe_tok_scales) // len M vector { static_assert(!with_group_zeropoints || with_group_scales); int M = A.size(0), N = B.size(1), K = A.size(1); - TORCH_CHECK(D.size(0) == M && D.size(1) == N); + STD_TORCH_CHECK(D.size(0) == M && D.size(1) == N); auto layout_A = make_cute_layout(A, "A"); auto layout_D = make_cute_layout(D, "D"); @@ -216,29 +218,29 @@ struct MacheteKernelTemplate { maybe_group_size == -1 ? K : maybe_group_size.value_or(K); int const scale_k = (K + group_size - 1) / group_size; - TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K); - TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N); + STD_TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K); + STD_TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N); if constexpr (with_group_scales) { - TORCH_CHECK(S_group_ptr && layout_S_group); - TORCH_CHECK((size<0>(*layout_S_group) == scale_k && - size<1>(*layout_S_group) == N)); + STD_TORCH_CHECK(S_group_ptr && layout_S_group); + STD_TORCH_CHECK((size<0>(*layout_S_group) == scale_k && + size<1>(*layout_S_group) == N)); } else { - TORCH_CHECK(!S_group_ptr, "Scales not supported"); + STD_TORCH_CHECK(!S_group_ptr, "Scales not supported"); } if constexpr (with_group_zeropoints) { - TORCH_CHECK(Z_group_ptr && layout_Z_group); - TORCH_CHECK((size<0>(*layout_Z_group) == scale_k && - size<1>(*layout_Z_group) == N)); - TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group, - "Scales and zeros must have the same layout"); + STD_TORCH_CHECK(Z_group_ptr && layout_Z_group); + STD_TORCH_CHECK((size<0>(*layout_Z_group) == scale_k && + size<1>(*layout_Z_group) == N)); + STD_TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group, + "Scales and zeros must have the same layout"); } else { - TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported"); + STD_TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported"); } if constexpr (with_channel_scales || with_token_scales) { - TORCH_CHECK( + STD_TORCH_CHECK( (maybe_ch_scales->numel() == N || maybe_ch_scales->numel() == 1) && (maybe_tok_scales->numel() == M || maybe_tok_scales->numel() == 1)); } @@ -298,11 +300,12 @@ struct MacheteKernelTemplate { Gemm gemm_op; cutlass::Status status = gemm_op.initialize(args, workspace, stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, - "Machete kernel failed to initialize workspace"); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Machete kernel failed to initialize workspace"); status = gemm_op.run(stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Machete kernel failed"); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Machete kernel failed"); } }; diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh b/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh new file mode 100644 index 00000000000..5ecf03daf79 --- /dev/null +++ b/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh @@ -0,0 +1,80 @@ +#pragma once + +#include "machete_mm_kernel.cuh" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "core/scalar_type.hpp" +#include "libtorch_stable/torch_utils.h" + +#include +#include + +#include +#include +#include + +namespace machete { + +struct MMArgs { + torch::stable::Tensor const& A; + torch::stable::Tensor const& B; + vllm::ScalarType const& b_type; + std::optional const& maybe_out_type; + std::optional const& maybe_group_scales; + std::optional const& maybe_group_zeros; + std::optional maybe_group_size; + std::optional const& maybe_channel_scales; + std::optional const& maybe_token_scales; + std::optional maybe_schedule; +}; + +struct SupportedSchedulesArgs { + torch::headeronly::ScalarType a_type; + vllm::ScalarType b_type; + std::optional maybe_group_scales_type; + std::optional maybe_group_zeros_type; + std::optional maybe_channel_scales_type; + std::optional maybe_token_scales_type; + std::optional maybe_out_type; +}; + +torch::stable::Tensor mm_dispatch(MMArgs args); + +std::vector supported_schedules_dispatch( + SupportedSchedulesArgs args); + +template +torch::stable::Tensor run_impl(MMArgs args) { + const torch::stable::accelerator::DeviceGuard device_guard( + args.A.get_device_index()); + + auto device = args.A.device(); + auto stream = get_current_cuda_stream(device.index()); + + int M = args.A.size(0); + int N = args.B.size(1); + int K = args.A.size(1); + + // Allocate output + torch::stable::Tensor D = torch::stable::empty( + {M, N}, equivalent_scalar_type_v, + std::nullopt, device); + + auto arguments = MacheteKernel::create_arguments( + stream, // + args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros, + args.maybe_group_size, args.maybe_channel_scales, + args.maybe_token_scales); + STD_TORCH_CHECK(MacheteKernel::can_implement(arguments), + "Machete kernel cannot be run with these arguments"); + + size_t workspace_size = MacheteKernel::get_workspace_size(arguments); + torch::stable::Tensor workspace = + torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, + std::nullopt, device); + + MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream); + + return D; +}; + +}; // namespace machete \ No newline at end of file diff --git a/csrc/quantization/machete/machete_prepack_kernel.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh similarity index 92% rename from csrc/quantization/machete/machete_prepack_kernel.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh index d002355ca49..6a28dceccc8 100644 --- a/csrc/quantization/machete/machete_prepack_kernel.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh @@ -2,7 +2,8 @@ #include "machete_mm_kernel.cuh" #include "cutlass_extensions/cute_utils.cuh" -#include "cutlass_extensions/torch_utils.hpp" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include namespace machete { @@ -60,8 +61,8 @@ static void prepack_B_template( auto ilvd_NKbNbKL_to_offset = PrepackedLayoutB::ilvd_NKbNbKL_to_offset(shape(B_layout)); - TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0); - TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0); + STD_TORCH_CHECK(size<0>(B_layout) % size<0>(TileShapeNKL{}) == 0); + STD_TORCH_CHECK(size<1>(B_layout) % size<1>(TileShapeNKL{}) == 0); auto N_tiles = size<0>(B_layout) / size<0>(TileShapeNKL{}); auto K_tiles = size<1>(B_layout) / size<1>(TileShapeNKL{}); diff --git a/csrc/quantization/machete/machete_prepack_launcher.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh similarity index 63% rename from csrc/quantization/machete/machete_prepack_launcher.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh index 634b651a4d1..63af5e496b8 100644 --- a/csrc/quantization/machete/machete_prepack_launcher.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh @@ -1,41 +1,49 @@ #pragma once #include "machete_prepack_kernel.cuh" -#include "cutlass_extensions/torch_utils.hpp" +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" +#include "libtorch_stable/torch_utils.h" + +#include +#include + +#include namespace machete { struct PrepackBArgs { - torch::Tensor const& B; - at::ScalarType a_type; + torch::stable::Tensor const& B; + torch::headeronly::ScalarType a_type; vllm::ScalarType b_type; - std::optional maybe_group_scales_type; + std::optional maybe_group_scales_type; }; template -torch::Tensor prepack_impl(torch::Tensor const B) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(B)); +torch::stable::Tensor prepack_impl(torch::stable::Tensor const& B) { + const torch::stable::accelerator::DeviceGuard device_guard( + B.get_device_index()); using ElementB = typename PrepackedLayoutB::ElementB; using PPBlockShape_NK = typename PrepackedLayoutB::PPBlockShape_NK; auto device = B.device(); - auto stream = at::cuda::getCurrentCUDAStream(device.index()); + auto stream = get_current_cuda_stream(device.index()); auto B_ptr = static_cast(B.const_data_ptr()); // elements per storage item for B auto eles_per_storage = - (B.dtype().itemsize() * 8) / cute::sizeof_bits_v; + (B.element_size() * 8) / cute::sizeof_bits_v; // torch B passed in is/should be (packed_K,N), the kernel expects (N,K,L) (to // match cutlass using (N,K,L) for B), so we transpose B to (N,packed_K,L) - auto Bt_packed = B.t(); + auto Bt_packed = torch::stable::transpose(B, 0, 1); - TORCH_CHECK( + STD_TORCH_CHECK( (B.size(0) * eles_per_storage) % size<1>(PPBlockShape_NK{}) == 0, "B.shape[0] (in terms of unpacked elements) must be a multiple of ", size<1>(PPBlockShape_NK{})); - TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0, - "B.shape[1] must be a multiple of ", size<0>(PPBlockShape_NK{})); + STD_TORCH_CHECK(B.size(1) % size<0>(PPBlockShape_NK{}) == 0, + "B.shape[1] must be a multiple of ", + size<0>(PPBlockShape_NK{})); using StrideB = cutlass::detail::TagToStrideB_t; auto const l_Bt_packed = make_cute_layout(Bt_packed, "B"); @@ -49,7 +57,7 @@ torch::Tensor prepack_impl(torch::Tensor const B) { // new_shape = (N, packed_K, L) * (1, eles_per_storage, 1) -> (N, K, L) // new_stride = (s0, s1, s2) * (eles_per_storage, 1, eles_per_storage) // when s1 == 1 - TORCH_CHECK(stride<1>(l_Bt_packed) == 1); + STD_TORCH_CHECK(stride<1>(l_Bt_packed) == 1); // clang-format off auto const layout_Bt = make_layout( transform_with_idx(l_Bt_packed.shape(), [&](auto ele, auto idx) { @@ -61,7 +69,9 @@ torch::Tensor prepack_impl(torch::Tensor const B) { // clang-format on // Allocate output - torch::Tensor D = torch::empty_like(B, {}, at::MemoryFormat::Contiguous); + torch::stable::Tensor D = torch::stable::empty( + B.sizes(), B.scalar_type(), std::nullopt, B.device(), std::nullopt, + torch::headeronly::MemoryFormat::Contiguous); prepack_B_template( stream, B_ptr, layout_Bt, static_cast(D.mutable_data_ptr())); @@ -69,6 +79,6 @@ torch::Tensor prepack_impl(torch::Tensor const B) { return D; }; -torch::Tensor prepack_B_dispatch(PrepackBArgs args); +torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args); }; // namespace machete \ No newline at end of file diff --git a/csrc/quantization/machete/machete_prepacked_layout.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh similarity index 99% rename from csrc/quantization/machete/machete_prepacked_layout.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh index 4a7d6341e6c..c16a2ab8a33 100644 --- a/csrc/quantization/machete/machete_prepacked_layout.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh @@ -1,9 +1,5 @@ #pragma once -#include -#include -#include - // clang-format off // The cutlass include order matters (annoyingly) diff --git a/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu b/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu new file mode 100644 index 00000000000..7736d5b3ece --- /dev/null +++ b/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu @@ -0,0 +1,77 @@ +#include "machete_mm_launcher.cuh" +#include "machete_prepack_launcher.cuh" +#include "core/scalar_type.hpp" + +#include +#include +#include + +namespace machete { + +using namespace vllm; + +std::vector supported_schedules( + torch::headeronly::ScalarType a_type, int64_t b_type_id, + std::optional maybe_group_scales_type, + std::optional maybe_group_zeros_type, + std::optional maybe_channel_scales_type, + std::optional maybe_token_scales_type, + std::optional maybe_out_type) { + ScalarType const b_type = ScalarType::from_id(b_type_id); + return supported_schedules_dispatch({ + .a_type = a_type, + .b_type = b_type, + .maybe_group_scales_type = maybe_group_scales_type, + .maybe_group_zeros_type = maybe_group_zeros_type, + .maybe_channel_scales_type = maybe_channel_scales_type, + .maybe_token_scales_type = maybe_token_scales_type, + .maybe_out_type = maybe_out_type, + }); +} + +torch::stable::Tensor mm( + torch::stable::Tensor const& A, torch::stable::Tensor const& B, + int64_t b_type_id, + std::optional const& maybe_out_type, + std::optional const& maybe_group_scales, + std::optional const& maybe_group_zeros, + std::optional maybe_group_size, + std::optional const& maybe_channel_scales, + std::optional const& maybe_token_scales, + std::optional maybe_schedule) { + ScalarType const b_type = ScalarType::from_id(b_type_id); + return mm_dispatch({.A = A, + .B = B, + .b_type = b_type, + .maybe_out_type = maybe_out_type, + .maybe_group_scales = maybe_group_scales, + .maybe_group_zeros = maybe_group_zeros, + .maybe_group_size = maybe_group_size, + .maybe_channel_scales = maybe_channel_scales, + .maybe_token_scales = maybe_token_scales, + .maybe_schedule = maybe_schedule}); +} + +torch::stable::Tensor prepack_B( + torch::stable::Tensor const& B, torch::headeronly::ScalarType const& a_type, + int64_t b_type_id, + std::optional const& + maybe_group_scales_type) { + ScalarType const b_type = ScalarType::from_id(b_type_id); + return prepack_B_dispatch( + {.B = B, + .a_type = a_type, + .b_type = b_type, + .maybe_group_scales_type = maybe_group_scales_type}); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("machete_prepack_B", TORCH_BOX(&prepack_B)); + m.impl("machete_mm", TORCH_BOX(&mm)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) { + m.impl("machete_supported_schedules", TORCH_BOX(&supported_schedules)); +} + +}; // namespace machete diff --git a/csrc/quantization/marlin/.gitignore b/csrc/libtorch_stable/quantization/marlin/.gitignore similarity index 100% rename from csrc/quantization/marlin/.gitignore rename to csrc/libtorch_stable/quantization/marlin/.gitignore diff --git a/csrc/quantization/marlin/awq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/awq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu index 307bae6738e..55ce5b4e732 100644 --- a/csrc/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" namespace marlin { @@ -218,56 +225,55 @@ __global__ void awq_marlin_repack_kernel( b_q_weight_ptr, out_ptr, size_k, size_n); \ } -torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, - int64_t size_n, int64_t num_bits, - bool is_a_8bit) { +torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight, + int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit) { // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, - " is not divisible by tile_k_size = ", marlin::tile_k_size); - TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, - " is not divisible by tile_n_size = ", marlin::tile_n_size); + STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", marlin::tile_k_size); + STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", marlin::tile_n_size); - TORCH_CHECK(num_bits == 4 || num_bits == 8, - "num_bits must be 4 or 8. Got = ", num_bits); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / num_bits; // Verify B - TORCH_CHECK(b_q_weight.size(0) == size_k, - "b_q_weight.size(0) = ", b_q_weight.size(0), - " is not size_k = ", size_k); - TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), - "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), - ", size_n = ", size_n, ", pack_factor = ", pack_factor); + STD_TORCH_CHECK(b_q_weight.size(0) == size_k, + "b_q_weight.size(0) = ", b_q_weight.size(0), + " is not size_k = ", size_k); + STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), + "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), + ", size_n = ", size_n, ", pack_factor = ", pack_factor); // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK( + b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, + "b_q_weight type is not kInt"); + + const int32_t device_index = b_q_weight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions() - .dtype(b_q_weight.dtype()) - .device(b_q_weight.device()); - torch::Tensor out = torch::empty( + torch::stable::Tensor out = torch::stable::empty( {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, - options); + b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); // Get ptrs uint32_t const* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); + STD_TORCH_CHECK(max_shared_mem > 0); if (false) { } @@ -276,13 +282,13 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, CALL_IF(4, true) CALL_IF(8, true) else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, - ", is_a_8bit = ", is_a_8bit); + STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", is_a_8bit = ", is_a_8bit); } return out; } -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("awq_marlin_repack", &awq_marlin_repack); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack)); } diff --git a/csrc/quantization/marlin/dequant.h b/csrc/libtorch_stable/quantization/marlin/dequant.h similarity index 100% rename from csrc/quantization/marlin/dequant.h rename to csrc/libtorch_stable/quantization/marlin/dequant.h diff --git a/csrc/quantization/marlin/generate_kernels.py b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py similarity index 99% rename from csrc/quantization/marlin/generate_kernels.py rename to csrc/libtorch_stable/quantization/marlin/generate_kernels.py index 7b316037ec6..2a038479893 100644 --- a/csrc/quantization/marlin/generate_kernels.py +++ b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py @@ -303,7 +303,7 @@ def generate_new_kernels(): if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: kernel_selector_str += ( "else if (a_type == vllm::kFE4M3fn)\n" - " TORCH_CHECK(false, " + " STD_TORCH_CHECK(false, " '"marlin kernel with fp8 activation is not built.");' ) diff --git a/csrc/quantization/marlin/gptq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/gptq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu index 796e6c5359d..cafa212bccb 100644 --- a/csrc/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" namespace marlin { @@ -275,64 +282,66 @@ __global__ void gptq_marlin_repack_kernel( b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ } -torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, - int64_t size_k, int64_t size_n, - int64_t num_bits, bool is_a_8bit) { +torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, + torch::stable::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit) { // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, - " is not divisible by tile_k_size = ", marlin::tile_k_size); - TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, - " is not divisible by tile_n_size = ", marlin::tile_n_size); + STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", marlin::tile_k_size); + STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", marlin::tile_n_size); - TORCH_CHECK(num_bits == 4 || num_bits == 8, - "num_bits must be 4 or 8. Got = ", num_bits); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / num_bits; // Verify B - TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), - ", size_k = ", size_k, ", pack_factor = ", pack_factor); - TORCH_CHECK(b_q_weight.size(1) == size_n, - "b_q_weight.size(1) = ", b_q_weight.size(1), - " is not size_n = ", size_n); + STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", pack_factor = ", pack_factor); + STD_TORCH_CHECK(b_q_weight.size(1) == size_n, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not size_n = ", size_n); // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK( + b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, + "b_q_weight type is not kInt"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); + STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); + STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int, + "perm type is not at::kInt"); + + const int32_t device_index = b_q_weight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions() - .dtype(b_q_weight.dtype()) - .device(b_q_weight.device()); - torch::Tensor out = torch::empty( + torch::stable::Tensor out = torch::stable::empty( {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, - options); + b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); // Detect if there is act_order bool has_perm = perm.size(0) != 0; // Get ptrs uint32_t const* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data_ptr()); - uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t const* perm_ptr = + reinterpret_cast(perm.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); + STD_TORCH_CHECK(max_shared_mem > 0); if (false) { } @@ -345,13 +354,13 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, CALL_IF(8, false, true) else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, - ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); + STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); } return out; } -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("gptq_marlin_repack", &gptq_marlin_repack); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack)); } diff --git a/csrc/quantization/marlin/kernel.h b/csrc/libtorch_stable/quantization/marlin/kernel.h similarity index 100% rename from csrc/quantization/marlin/kernel.h rename to csrc/libtorch_stable/quantization/marlin/kernel.h diff --git a/csrc/quantization/marlin/marlin.cu b/csrc/libtorch_stable/quantization/marlin/marlin.cu similarity index 61% rename from csrc/quantization/marlin/marlin.cu rename to csrc/libtorch_stable/quantization/marlin/marlin.cu index 721c206c33f..63fea239e4a 100644 --- a/csrc/quantization/marlin/marlin.cu +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cu @@ -24,7 +24,15 @@ #endif #include "kernel.h" -#include "core/registration.h" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ @@ -46,19 +54,22 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, } // namespace marlin -torch::Tensor marlin_gemm( - torch::Tensor& a, std::optional c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::Tensor& workspace, - vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, - int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, - bool is_zp_float) { - TORCH_CHECK_NOT_IMPLEMENTED(false, - "marlin_gemm(..) requires CUDA_ARCH >= 7.5"); - return torch::empty({1, 1}); +torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional const& perm_or_none, + torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, + int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, + bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { + STD_TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 7.5"); + return torch::stable::empty({1, 1}); } #else @@ -323,18 +334,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int thread_n_init, int sms, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { bool is_a_8bit = a_type.size_bits() == 8; - TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, - ", ", prob_n, ", ", prob_k, "]"); + STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); int group_blocks = 0; if (has_act_order) { if (is_k_full) { - TORCH_CHECK(group_size != -1); + STD_TORCH_CHECK(group_size != -1); group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } else { - TORCH_CHECK(group_size == 0); + STD_TORCH_CHECK(group_size == 0); group_blocks = 0; } } else { @@ -342,8 +353,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, group_blocks = -1; } else { group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } } @@ -384,25 +395,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + STD_TORCH_CHECK(max_shared_mem > 0); int major_capability, minor_capability; cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, dev); cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, dev); - TORCH_CHECK(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); int stages = 4; if (major_capability == 7 && minor_capability == 5) { stages = 2; - TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); + STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); } if (a_type == vllm::kFE4M3fn) { - TORCH_CHECK(major_capability * 10 + minor_capability >= 89, - "FP8 only support Ada Lovelace or newer GPUs."); - TORCH_CHECK( + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); + STD_TORCH_CHECK( major_capability * 10 + minor_capability == 89 || major_capability == 12, "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " @@ -432,10 +443,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, if (thread_k != -1 && thread_n != -1) { thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; exec_cfg = exec_config_t{1, thread_tfg}; - TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); - TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); + STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); } else { // Auto config exec_cfg = determine_exec_config( @@ -474,7 +485,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int thread_k_blocks = thread_k / 16; int thread_n_blocks = thread_n / 16; - TORCH_CHECK( + STD_TORCH_CHECK( is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages, @@ -495,14 +506,15 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, num_threads, is_zp_float, stages); if (kernel == MarlinDefault) { - TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, - ", ", prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", prob_m_split = ", prob_m_split, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, - ", num_threads = ", num_threads, ", num_bits = ", num_bits); + STD_TORCH_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, + ", prob_m_split = ", prob_m_split, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, + ", num_threads = ", num_threads, ", num_bits = ", num_bits); } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -530,71 +542,76 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, } // namespace marlin -torch::Tensor marlin_gemm( - torch::Tensor& a, std::optional c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& a_scales_or_none, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::Tensor& workspace, - vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, - int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, - bool is_zp_float) { +torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional const& perm_or_none, + torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, + int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, + bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; - auto c_dtype = a.dtype(); - if (a.scalar_type() == at::ScalarType::Half) { + auto c_scalar_type = a.scalar_type(); + if (a.scalar_type() == torch::headeronly::ScalarType::Half) { a_type_id = vllm::kFloat16.id(); c_type_id = vllm::kFloat16.id(); - } else if (a.scalar_type() == at::ScalarType::BFloat16) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { a_type_id = vllm::kBFloat16.id(); c_type_id = vllm::kBFloat16.id(); } else { - c_dtype = b_scales.dtype(); - if (b_scales.scalar_type() == at::ScalarType::Half) { + c_scalar_type = b_scales.scalar_type(); + if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { c_type_id = vllm::kBFloat16.id(); - TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); - torch::Tensor c = c_or_none.value(); - c_dtype = c.dtype(); + STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); + torch::stable::Tensor c = c_or_none.value(); + c_scalar_type = c.scalar_type(); - if (c.scalar_type() == at::ScalarType::Half) { + if (c.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (c.scalar_type() == at::ScalarType::BFloat16) { + } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { - TORCH_CHECK(false, "unsupported c dtype"); + STD_TORCH_CHECK(false, "unsupported c dtype"); } } - if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) { a_type_id = vllm::kFE4M3fn.id(); - } else if (a.scalar_type() == at::ScalarType::Char) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) { a_type_id = vllm::kS8.id(); } else { - TORCH_CHECK(false, "unsupported `a` scalar_type"); + STD_TORCH_CHECK(false, "unsupported `a` scalar_type"); } } s_type_id = c_type_id; if (b_type_id == vllm::kFE2M1f.id()) { - if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn) { s_type_id = vllm::kFE4M3fn.id(); - } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } else { - TORCH_CHECK(false, - "When b_type = float4_e2m1f, b_scale scalar type must be", - "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); + STD_TORCH_CHECK( + false, "When b_type = float4_e2m1f, b_scale scalar type must be", + "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); } } else if (b_type_id == vllm::kFE4M3fn.id() && - b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } @@ -606,54 +623,58 @@ torch::Tensor marlin_gemm( int pack_factor = 32 / b_type.size_bits(); // Verify A - TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), - ", size_m = ", size_m); - TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), - ", size_k = ", size_k); + STD_TORCH_CHECK(a.size(0) == size_m, + "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + STD_TORCH_CHECK(a.size(1) == size_k, + "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); // Verify B - TORCH_CHECK( + STD_TORCH_CHECK( size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k, " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), - ", size_k = ", size_k, - ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK( + STD_TORCH_CHECK( + (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, + ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); + STD_TORCH_CHECK( b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0, "b_q_weight.size(1) = ", b_q_weight.size(1), " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); int actual_size_n = (b_q_weight.size(1) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor; - TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, - ", actual_size_n = ", actual_size_n); + STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); // Verify device and strides - TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); - TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1"); + STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + STD_TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1"); // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes - TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8"); - TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes"); + STD_TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8"); + STD_TORCH_CHECK(((uint64_t)a.const_data_ptr()) % 16 == 0, + "A must aligned to 16 bytes"); - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); - TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); - torch::Tensor a_scales; - auto options = torch::TensorOptions().dtype(c_dtype).device(a.device()); - auto options_fp32 = - torch::TensorOptions().dtype(at::kFloat).device(a.device()); + torch::stable::Tensor a_scales; + const auto device = a.device(); if (a_scales_or_none.has_value()) { a_scales = a_scales_or_none.value(); - TORCH_CHECK(a_type.size_bits() == 8, - "a_scales can only be used for 8bit activation."); + STD_TORCH_CHECK(a_type.size_bits() == 8, + "a_scales can only be used for 8bit activation."); } else { - a_scales = torch::empty({0}, options_fp32); - TORCH_CHECK(a_type.size_bits() != 8, - "the a_scales parameter must be passed for 8bit activation."); + a_scales = torch::stable::empty({0}, torch::headeronly::ScalarType::Float, + std::nullopt, device); + STD_TORCH_CHECK( + a_type.size_bits() != 8, + "the a_scales parameter must be passed for 8bit activation."); } // thread_k: `k` size of a thread_tile in `weights` (can usually be left as @@ -664,84 +685,93 @@ torch::Tensor marlin_gemm( int thread_n = -1; // sms: number of SMs to use for the kernel int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device()); + const int32_t device_index = a.get_device_index(); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - torch::Tensor c; + torch::stable::accelerator::DeviceGuard device_guard(device_index); + torch::stable::Tensor c; if (c_or_none.has_value()) { c = c_or_none.value(); - TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); - TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); - TORCH_CHECK(c.size(0) == size_m, "Shape mismatch: c.size(0) = ", c.size(0), - ", size_m = ", size_m); - TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1), - ", size_n = ", size_n); + STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); + STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); + STD_TORCH_CHECK(c.size(0) == size_m, + "Shape mismatch: c.size(0) = ", c.size(0), + ", size_m = ", size_m); + STD_TORCH_CHECK(c.size(1) == size_n, + "Shape mismatch: c.size(1) = ", c.size(1), + ", size_n = ", size_n); } else { - c = torch::empty({size_m, size_n}, options); + c = torch::stable::empty({size_m, size_n}, c_scalar_type, std::nullopt, + device); } if (size_m == 0) return c; // Alloc C tmp buffer that is going to be used for the global reduce - torch::Tensor c_tmp; + torch::stable::Tensor c_tmp; if (use_fp32_reduce) { int max_m_block_size = (size_m + 16 - 1) / 16 * 16; max_m_block_size = min(max_m_block_size, 64); int max_c_tmp_size = sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n; - c_tmp = torch::empty({max_c_tmp_size}, options_fp32); + c_tmp = torch::stable::empty({max_c_tmp_size}, + torch::headeronly::ScalarType::Float, + std::nullopt, device); } else { - c_tmp = torch::empty({0}, options_fp32); + c_tmp = torch::stable::empty({0}, torch::headeronly::ScalarType::Float, + std::nullopt, device); } // Detect groupsize and act_order int num_groups = -1; int group_size = -1; - int rank = b_scales.sizes().size(); - TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2"); - TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1), - " is not size_n = ", size_n); + int rank = b_scales.dim(); + STD_TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2"); + STD_TORCH_CHECK(b_scales.size(1) == size_n, + "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); num_groups = b_scales.size(0); - torch::Tensor g_idx, perm, a_tmp; + torch::stable::Tensor g_idx, perm, a_tmp; if (g_idx_or_none.has_value() && perm_or_none.has_value()) { g_idx = g_idx_or_none.value(); perm = perm_or_none.value(); - TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); // Verify g_idx and perm - TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || - (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", g_idx.size(-1), - " and perm.size(-1) = ", perm.size(-1), - ", where size_k = ", size_k); + STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || + (g_idx.size(-1) == size_k && perm.size(-1) == size_k), + "Unexpected g_idx.size(-1) = ", g_idx.size(-1), + " and perm.size(-1) = ", perm.size(-1), + ", where size_k = ", size_k); } else { - g_idx = torch::empty({0}, options); - perm = torch::empty({0}, options); - a_tmp = torch::empty({0}, options); + g_idx = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); + perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); + a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; if (has_act_order) { - a_tmp = torch::empty({size_m, size_k}, options); + a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt, + device); if (is_k_full) { - TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); + STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); group_size = size_k / num_groups; } else { group_size = 0; } } else { - a_tmp = torch::empty({0}, options); + a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); if (num_groups > 1) { - TORCH_CHECK( + STD_TORCH_CHECK( size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); group_size = size_k / num_groups; @@ -750,109 +780,114 @@ torch::Tensor marlin_gemm( } } - torch::Tensor global_scale; + torch::stable::Tensor global_scale; if (global_scale_or_none.has_value()) { global_scale = global_scale_or_none.value(); - TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "global_scale can only be used for nvfp4 format."); + STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, + "global_scale can only be used for nvfp4 format."); } else { - global_scale = torch::empty({0}, options_fp32); - TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), - "the global_scale parameter must be passed for nvfp4 format."); + global_scale = torch::stable::empty( + {0}, torch::headeronly::ScalarType::Float, std::nullopt, device); + STD_TORCH_CHECK( + !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), + "the global_scale parameter must be passed for nvfp4 format."); } bool has_bias = b_bias_or_none.has_value(); - torch::Tensor b_bias; + torch::stable::Tensor b_bias; if (has_bias) { b_bias = b_bias_or_none.value(); - TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); - TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); - TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n"); - TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1"); + STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); + STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); + STD_TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n"); + STD_TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1"); } else { - b_bias = torch::empty({0}, options); + b_bias = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } - torch::Tensor b_zeros; + torch::stable::Tensor b_zeros; if (b_zeros_or_none.has_value()) { b_zeros = b_zeros_or_none.value(); - TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); - TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); + STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); + STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); } else { - b_zeros = torch::empty({0}, options); + b_zeros = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } bool has_zp = b_zeros.size(-1) > 0; if (has_zp) { - TORCH_CHECK( + STD_TORCH_CHECK( b_type == vllm::kU4 || b_type == vllm::kU8, "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str()); } else { - TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || - b_type == vllm::kS4 || b_type == vllm::kS8 || - b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, - "b_type must be uint4b8, uint8b128, int4, int8, " - "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", - b_type.str()); + STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || + b_type == vllm::kS4 || b_type == vllm::kS8 || + b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, + "b_type must be uint4b8, uint8b128, int4, int8, " + "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", + b_type.str()); } if (has_zp && is_zp_float) { - TORCH_CHECK(a.scalar_type() == at::ScalarType::Half, - "Computation type must be float16 (half) when using float zero " - "points."); + STD_TORCH_CHECK( + a.scalar_type() == torch::headeronly::ScalarType::Half, + "Computation type must be float16 (half) when using float zero " + "points."); } // Verify b_zeros if (has_zp) { - int rank = b_zeros.sizes().size(); - TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2"); + int rank = b_zeros.dim(); + STD_TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2"); if (is_zp_float) { - TORCH_CHECK(b_zeros.size(1) == size_n, - "b_zeros dim 1 = ", b_zeros.size(1), - " is not size_n = ", size_n); - TORCH_CHECK(num_groups == b_zeros.size(0), - "b_zeros dim 0 = ", b_zeros.size(0), - " is not num_groups = ", num_groups); - TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); + STD_TORCH_CHECK(b_zeros.size(1) == size_n, + "b_zeros dim 1 = ", b_zeros.size(1), + " is not size_n = ", size_n); + STD_TORCH_CHECK(num_groups == b_zeros.size(0), + "b_zeros dim 0 = ", b_zeros.size(0), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); } else { - TORCH_CHECK(b_zeros.size(0) == num_groups, - "b_zeros dim 0 = ", b_zeros.size(0), - " is not num_groups = ", num_groups); - TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor, - "b_zeros dim 1 = ", b_zeros.size(1), - " is not size_n / pack_factor = ", size_n / pack_factor); + STD_TORCH_CHECK(b_zeros.size(0) == num_groups, + "b_zeros dim 0 = ", b_zeros.size(0), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor, + "b_zeros dim 1 = ", b_zeros.size(1), + " is not size_n / pack_factor = ", size_n / pack_factor); } } // Verify workspace size - TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, - "size_n = ", size_n, ", is not divisible by min_thread_n = ", - MARLIN_NAMESPACE_NAME::min_thread_n); + STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, + "size_n = ", size_n, ", is not divisible by min_thread_n = ", + MARLIN_NAMESPACE_NAME::min_thread_n); int min_workspace_size = sms; - TORCH_CHECK(workspace.numel() >= min_workspace_size, - "workspace.numel = ", workspace.numel(), - " is below min_workspace_size = ", min_workspace_size); + STD_TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); - int dev = a.get_device(); - - TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float, - "scalar type of a_scales must be float"); - TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, - "scalar type of global_scale must be float"); + STD_TORCH_CHECK( + a_scales.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of a_scales must be float"); + STD_TORCH_CHECK( + global_scale.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of global_scale must be float"); if (a_type.size_bits() == 16) { - TORCH_CHECK( + STD_TORCH_CHECK( a.scalar_type() == c.scalar_type(), "scalar type of a must be the same with c for 16 bit activation"); } marlin::marlin_mm( - a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), - b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(), - global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), - perm.data_ptr(), a_tmp.data_ptr(), size_m, size_n, size_k, a.stride(0), - workspace.data_ptr(), a_type, b_type, c_type, s_type, has_bias, - has_act_order, is_k_full, has_zp, num_groups, group_size, dev, - at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(), + c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), + a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), + global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), + g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), + a_tmp.mutable_data_ptr(), size_m, size_n, size_k, a.stride(0), + workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, has_bias, + has_act_order, is_k_full, has_zp, num_groups, group_size, device_index, + get_current_cuda_stream(device_index), thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float); return c; @@ -860,6 +895,6 @@ torch::Tensor marlin_gemm( #endif -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("marlin_gemm", &marlin_gemm); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("marlin_gemm", TORCH_BOX(&marlin_gemm)); } diff --git a/csrc/quantization/marlin/marlin.cuh b/csrc/libtorch_stable/quantization/marlin/marlin.cuh similarity index 93% rename from csrc/quantization/marlin/marlin.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin.cuh index d3a91568349..bfb65e874b3 100644 --- a/csrc/quantization/marlin/marlin.cuh +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cuh @@ -2,14 +2,6 @@ #ifndef _marlin_cuh #define _marlin_cuh - // These torch headers are only needed by non-stable callers (e.g. ops.cu). - // Guard them so that stable ABI targets can still include marlin.cuh - // for Vec, constants, and cp_async helpers without pulling in torch/all.h. - #ifndef TORCH_TARGET_VERSION - #include - #include - #include - #endif #include #include #include diff --git a/csrc/quantization/marlin/marlin_dtypes.cuh b/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/quantization/marlin/marlin_dtypes.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu new file mode 100644 index 00000000000..f8ef6b12a01 --- /dev/null +++ b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -0,0 +1,118 @@ + +#include "marlin.cuh" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" + +// for only non-zp format (like gptq) +__global__ void marlin_int4_fp8_preprocess_kernel_without_zp( + // qweight: (size_k * size_n // 8,) + const int32_t* __restrict__ qweight, + // output: same shape with qweight + int32_t* __restrict__ output) { + int32_t val = qweight[blockIdx.x * 32 + threadIdx.x]; + int32_t new_val = 0; + +#pragma unroll + for (int32_t i = 0; i < 8; i++) { + int32_t single_val = val & 0xF; + single_val = single_val >= 8 ? single_val - 8 : 15 - single_val; + new_val |= single_val << (i * 4); + val >>= 4; + } + + output[blockIdx.x * 32 + threadIdx.x] = new_val; +} + +// for awq format only (with zp and with awq weight layout) +__global__ void marlin_int4_fp8_preprocess_kernel_awq( + // AWQ qweight: (size_k, size_n // 8) + const int32_t* __restrict__ qweight, + // output: same shape with qweight + int32_t* __restrict__ output, + // AWQ zeros: (size_k // group_size, size_n // 8) + const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, + int32_t group_size) { + int32_t val = + qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; + int32_t zero = + qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + + blockIdx.y]; + int32_t new_val = 0; + +#pragma unroll + for (int32_t i = 0; i < 8; i++) { + int32_t single_val = val & 0xF; + int32_t single_zero = zero & 0xF; + + single_val = + single_val >= single_zero ? single_val - single_zero : 15 - single_val; + new_val |= single_val << (i * 4); + val >>= 4; + zero >>= 4; + } + + output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; +} + +torch::stable::Tensor marlin_int4_fp8_preprocess( + torch::stable::Tensor& qweight, + std::optional qzeros_or_none, bool inplace) { + STD_TORCH_CHECK(qweight.is_cuda(), "qweight is not on GPU"); + STD_TORCH_CHECK(qweight.scalar_type() == torch::headeronly::ScalarType::Int, + "qweight.dtype != torch.int32"); + + const int32_t device_index = qweight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor output = + inplace ? qweight : torch::stable::empty_like(qweight); + + if (!qzeros_or_none.has_value()) { + STD_TORCH_CHECK(qweight.numel() * 8 % 256 == 0, + "qweight.numel() * 8 % 256 != 0"); + + int blocks = qweight.numel() * 8 / 256; + marlin_int4_fp8_preprocess_kernel_without_zp<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr())); + } else { + int32_t size_k = qweight.size(0); + int32_t size_n = qweight.size(1) * 8; + torch::stable::Tensor qzeros = qzeros_or_none.value(); + + STD_TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0"); + STD_TORCH_CHECK(qzeros.is_cuda(), "qzeros is not on GPU"); + STD_TORCH_CHECK(qzeros.scalar_type() == torch::headeronly::ScalarType::Int, + "qweight.dtype != torch.int32"); + STD_TORCH_CHECK(qzeros.get_device_index() == device_index, + "qzeros is not on the same device with qweight"); + + int32_t group_size = qweight.size(0) / qzeros.size(0); + STD_TORCH_CHECK(qweight.size(1) == qzeros.size(1), + "qweight.size(1) != qzeros.size(1)"); + STD_TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0, + "qweight.size(0) % qzeros.size(0) != 0"); + STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); + + dim3 blocks(size_k / 32, size_n / 8); + marlin_int4_fp8_preprocess_kernel_awq<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(qzeros.const_data_ptr()), size_n, + size_k, group_size); + } + + return output; +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("marlin_int4_fp8_preprocess", TORCH_BOX(&marlin_int4_fp8_preprocess)); +} diff --git a/csrc/quantization/marlin/marlin_mma.h b/csrc/libtorch_stable/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/quantization/marlin/marlin_mma.h rename to csrc/libtorch_stable/quantization/marlin/marlin_mma.h diff --git a/csrc/quantization/marlin/marlin_template.h b/csrc/libtorch_stable/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/quantization/marlin/marlin_template.h rename to csrc/libtorch_stable/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/libtorch_stable/quantization/vectorization_utils.cuh index 98b491b7e23..0cc89bf289d 100644 --- a/csrc/libtorch_stable/quantization/vectorization_utils.cuh +++ b/csrc/libtorch_stable/quantization/vectorization_utils.cuh @@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment( ScaOp&& scalar_op) { // InT -> OutT static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, "VEC_SIZE must be a positive power-of-two"); - constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B + constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 16 B + constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT); // eg: 16 B uintptr_t addr = reinterpret_cast(in); + uintptr_t out_addr = reinterpret_cast(out); - // fast path when the whole region is already aligned - // Note: currently the output is guaranteed to be same as the input, so we - // don't check it here, comments here just for future reference. - bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); + // fast path when input and output are both fully aligned. The vector + // load/store below go through vec_n_t, declared + // __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its + // own vector width. out is NOT generally co-aligned with in: e.g. + // reshape_and_cache_flash writes KV-cache rows whose byte offset is a + // multiple of head_size, which for head sizes that are not a multiple + // of VEC_SIZE puts some rows off the vector-width boundary. + bool can_vec = ((addr & (WIDTH - 1)) == 0) && + ((out_addr & (OUT_WIDTH - 1)) == 0) && + ((len & (VEC_SIZE - 1)) == 0); if (can_vec) { int num_vec = len / VEC_SIZE; @@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment( prefix_elems /= sizeof(InT); prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16 + // the prefix below aligns in; if that does not also align out (their + // addresses differ modulo the vector width), vectorizing is impossible + // and the whole copy must stay scalar. + if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) { + for (int i = tid; i < len; i += stride) { + scalar_op(out[i], in[i]); + } + return; + } + // 1. prefill the when it is unsafe to vectorize for (int i = tid; i < prefix_elems; i += stride) { scalar_op(out[i], in[i]); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index 1eed7579924..a93a136bf42 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -19,7 +19,7 @@ #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/util/packed_stride.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on @@ -37,6 +37,7 @@ void cutlass_gemm_caller( typename GemmKernel::MainloopArguments mainloop_args, typename GemmKernel::EpilogueArguments epilogue_args, typename GemmKernel::TileSchedulerArguments scheduler = {}) { + const torch::stable::accelerator::DeviceGuard device_guard(device.index()); cutlass::KernelHardwareInfo hw_info; typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm, prob_shape, diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index 4cb591be056..7b7d4d71473 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -14,7 +14,7 @@ #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu index bc088cf633f..5178eed0722 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu index f3df69850ec..b7930012265 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu index 7ceb0697df2..426025ac131 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu index d3318c48767..a97909d37ae 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh index cf62e81fd75..529b28ceece 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh @@ -25,33 +25,43 @@ using namespace cute; template + class EpilogueScheduler, class MainloopScheduler, + bool swap_ab_ = false> struct cutlass_3x_gemm_fp8_blockwise { + static constexpr bool swap_ab = swap_ab_; using ElementAB = cutlass::float_e4m3_t; using ElementA = ElementAB; using LayoutA = cutlass::layout::RowMajor; + using LayoutA_Transpose = typename cutlass::layout::LayoutTranspose::type; static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; using ElementB = ElementAB; using LayoutB = cutlass::layout::ColumnMajor; + using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose::type; static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; using ElementD = OutType; using LayoutD = cutlass::layout::RowMajor; + using LayoutD_Transpose = typename cutlass::layout::LayoutTranspose::type; static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; using ElementC = void; // TODO: support bias using LayoutC = LayoutD; + using LayoutC_Transpose = LayoutD_Transpose; static constexpr int AlignmentC = AlignmentD; using ElementAccumulator = float; using ElementCompute = float; using ElementBlockScale = float; - using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig< + using ScaleConfig = conditional_t; + cute::GMMA::Major::K, cute::GMMA::Major::MN>, + cutlass::detail::Sm90BlockwiseScaleConfig< + ScaleGranularityM, ScaleGranularityN, ScaleGranularityK, + cute::GMMA::Major::MN, cute::GMMA::Major::K>>; using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); @@ -71,30 +81,46 @@ struct cutlass_3x_gemm_fp8_blockwise { ElementAccumulator, ElementCompute, ElementC, - LayoutC, + conditional_t, AlignmentC, ElementD, - LayoutD, + conditional_t, AlignmentD, EpilogueScheduler, DefaultOperation >::CollectiveOp; - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, - MainloopScheduler - >::CollectiveOp; + using CollectiveMainloop = conditional_t, + AlignmentB, + ElementA, + cute::tuple, + AlignmentA, + ElementAccumulator, + MmaTileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopScheduler + >::CollectiveOp, + typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + MmaTileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopScheduler + >::CollectiveOp>; using KernelType = enable_sm90_or_later, CollectiveMainloop, CollectiveEpilogue>>; @@ -107,6 +133,7 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales, torch::stable::Tensor const& b_scales) { + static constexpr bool swap_ab = Gemm::swap_ab; using GemmKernel = typename Gemm::GemmKernel; using StrideA = typename Gemm::GemmKernel::StrideA; using StrideB = typename Gemm::GemmKernel::StrideB; @@ -122,8 +149,6 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te int32_t m = a.size(0), n = b.size(1), k = a.size(1); - STD_TORCH_CHECK(m % 4 == 0, "m must be divisible by 4"); - StrideA a_stride; StrideB b_stride; StrideC c_stride; @@ -132,12 +157,16 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); c_stride = - cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + cutlass::make_cute_packed_stride( + StrideC{}, swap_ab ? cute::make_shape(n, m, 1) + : cute::make_shape(m, n, 1)); - LayoutSFA layout_SFA = - ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_SFB = - ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + LayoutSFA layout_SFA = swap_ab + ? ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) + : ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + LayoutSFB layout_SFB = swap_ab + ? ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) + : ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); auto a_ptr = static_cast(a.data_ptr()); auto b_ptr = static_cast(b.data_ptr()); @@ -145,15 +174,25 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te auto b_scales_ptr = static_cast(b_scales.data_ptr()); typename GemmKernel::MainloopArguments mainloop_args{}; - mainloop_args.ptr_A = a_ptr; - mainloop_args.dA = a_stride; - mainloop_args.ptr_B = b_ptr; - mainloop_args.dB = b_stride; - mainloop_args.ptr_SFA = a_scales_ptr; mainloop_args.layout_SFA = layout_SFA; - mainloop_args.ptr_SFB = b_scales_ptr; mainloop_args.layout_SFB = layout_SFB; - auto prob_shape = cute::make_shape(m, n, k, 1); + if (swap_ab) { + mainloop_args.ptr_A = b_ptr; + mainloop_args.dA = b_stride; + mainloop_args.ptr_B = a_ptr; + mainloop_args.dB = a_stride; + mainloop_args.ptr_SFA = b_scales_ptr; + mainloop_args.ptr_SFB = a_scales_ptr; + } else { + mainloop_args.ptr_A = a_ptr; + mainloop_args.dA = a_stride; + mainloop_args.ptr_B = b_ptr; + mainloop_args.dB = b_stride; + mainloop_args.ptr_SFA = a_scales_ptr; + mainloop_args.ptr_SFB = b_scales_ptr; + } + auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) + : cute::make_shape(m, n, k, 1); auto c_ptr = static_cast(out.data_ptr()); typename GemmKernel::EpilogueArguments epilogue_args{ @@ -168,12 +207,21 @@ void cutlass_gemm_blockwise_sm90_fp8_dispatch(torch::stable::Tensor& out, torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales, torch::stable::Tensor const& b_scales) { - // TODO: better heuristics + bool swap_ab = (a.size(0) % 4) != 0; + if (!swap_ab) { + cutlass_gemm_caller_blockwise, + Shape<_1, _2, _1>, cutlass::epilogue::TmaWarpSpecializedCooperative, + cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum>>( + out, a, b, a_scales, b_scales); + return; + } + cutlass_gemm_caller_blockwise, - Shape<_1, _2, _1>, cutlass::epilogue::TmaWarpSpecializedCooperative, - cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum>>( - out, a, b, a_scales, b_scales); + OutType, 128, 1, 128, Shape<_128, _16, _128>, + Shape<_1, _1, _1>, cutlass::epilogue::TmaWarpSpecialized, + cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8BlockScaledAccum, + true>>(out, a, b, a_scales, b_scales); } } // namespace vllm \ No newline at end of file diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh index f790b3653d5..42f9e4d5bc8 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM100 (fp8) based on the diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu index 972d6c62606..1f709699eb9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm120_fp8_dispatch.cuh" #include "core/batch_invariant.hpp" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh index f78b8daea51..2fae3016c30 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu index 717a2a58830..bf2bd030e3c 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh index e073b4e64ad..6942c4b48cc 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh @@ -74,6 +74,8 @@ void run_get_group_gemm_starts( bool per_act_token = a_scales.numel() != 1; bool per_out_ch = b_scales.numel() != num_experts; + const torch::stable::accelerator::DeviceGuard device_guard( + a_tensors.get_device_index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index()); if (false) { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh index b523d7baeaa..b4cd520e967 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh @@ -7,8 +7,9 @@ #include "cutlass/gemm/device/gemm_universal_adapter.h" #include -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" +#include "libtorch_stable/torch_utils.h" #include "get_group_starts.cuh" using namespace cute; @@ -103,6 +104,8 @@ void cutlass_group_gemm_caller(torch::stable::Tensor& out_tensors, int num_experts = static_cast(expert_offsets.size(0)); + const torch::stable::accelerator::DeviceGuard device_guard( + a_tensors.get_device_index()); auto stream = get_current_cuda_stream(a_tensors.get_device_index()); auto device = a_tensors.device(); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu index 2632989cc69..47e0985b123 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu @@ -212,6 +212,8 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller( "n and k must fit in int32"); int const num_experts = static_cast(num_experts64); + const torch::stable::accelerator::DeviceGuard device_guard( + expert_first_token_offset.get_device_index()); auto stream = get_current_cuda_stream(expert_first_token_offset.get_device_index()); @@ -241,6 +243,7 @@ void get_cutlass_moe_mm_data_caller( const std::optional& blockscale_offsets, const bool is_gated) { auto device = topk_ids.device(); + const torch::stable::accelerator::DeviceGuard device_guard(device.index()); auto stream = get_current_cuda_stream(device.index()); torch::stable::Tensor atomic_buffer = torch::stable::new_zeros( topk_ids, {num_experts}, torch::headeronly::ScalarType::Int); @@ -311,6 +314,8 @@ void get_cutlass_batched_moe_mm_data_caller( const torch::stable::Tensor& expert_num_tokens, const int64_t num_local_experts, const int64_t padded_m, const int64_t n, const int64_t k) { + const torch::stable::accelerator::DeviceGuard device_guard( + expert_offsets.get_device_index()); auto stream = get_current_cuda_stream(expert_offsets.get_device_index()); if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 7846e609fe7..9f9e12e982e 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -22,7 +22,7 @@ #include "cutlass/epilogue/threadblock/fusion/visitors.hpp" #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" -#include "core/math.hpp" +#include "libtorch_stable/core/math.hpp" #include "libtorch_stable/cutlass_extensions/common.hpp" // clang-format on @@ -156,6 +156,7 @@ inline void cutlass_gemm_caller(torch::stable::Tensor& out, torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt, device); + const torch::stable::accelerator::DeviceGuard device_guard(device.index()); auto stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemm_op.can_implement(args)); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu index 0f9873cbf88..51f84d2ffd9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -1,3 +1,4 @@ +#include #include #include @@ -174,15 +175,20 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) { bool cutlass_group_gemm_supported(int64_t cuda_device_capability) { // CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper) - // or CUDA 12.8 and SM100 (Blackwell) + // or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an + // actual cutlass_moe_mm dispatch compiled into this file. #if defined CUDA_VERSION - if (cuda_device_capability >= 100) { + #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100 + if (cuda_device_capability >= 100 && cuda_device_capability < 120) { return CUDA_VERSION >= 12080; } - if (cuda_device_capability >= 90) { + #endif + #if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90 + if (cuda_device_capability >= 90 && cuda_device_capability < 100) { return CUDA_VERSION >= 12030; } + #endif #endif return false; diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu index d02fc2296e6..86696e2f7a7 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu @@ -1,6 +1,6 @@ #include "../../../../quantization/w8a8/fp8/common.cuh" #include "../../../dispatch_utils.h" -#include "../../../../cub_helpers.h" +#include "../../../cub_helpers.h" #include "../../vectorization_utils.cuh" #include "../../../torch_utils.h" #include 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 316a7d37522..f4a0bd54428 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 @@ -119,6 +119,10 @@ __global__ void per_token_group_quant_8bit_kernel( static_cast(output_q) + block_group_offset; scale_element_t* scale_output; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + if constexpr (IS_COLUMN_MAJOR) { const int num_elems_per_pack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); @@ -153,6 +157,10 @@ __global__ void per_token_group_quant_8bit_kernel( QuantizeGroup(smem_group, group_output, group_size, lane_id, threads_per_group, y_s, min_8bit, max_8bit); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif } inline int GetGroupsPerBlock(int64_t num_groups) { @@ -195,6 +203,8 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, STD_TORCH_CHECK(input.numel() % group_size == 0); STD_TORCH_CHECK(output_s.dim() == 2); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); cudaStream_t stream = get_current_cuda_stream(); constexpr int THREADS_PER_GROUP = 16; @@ -209,45 +219,56 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, const int scale_num_rows = output_s.size(1); const int scale_stride = output_s.stride(1); -#define LAUNCH_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); \ - if (is_column_major) { \ - if (scale_ue8m0) { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit, scale_num_rows, scale_stride); \ - } else { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit, scale_num_rows, scale_stride); \ - } \ - } else { \ - if (scale_ue8m0) { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit); \ - } else { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit); \ - } \ - } \ +#ifndef USE_ROCM + #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ + do { \ + cudaLaunchConfig_t config = {}; \ + config.gridDim = dim3(num_blocks); \ + config.blockDim = dim3(num_threads); \ + config.dynamicSmemBytes = (SMEM_BYTES); \ + config.stream = stream; \ + cudaLaunchAttribute attrs[1]; \ + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ + attrs[0].val.programmaticStreamSerializationAllowed = 1; \ + config.numAttrs = 1; \ + config.attrs = attrs; \ + cudaLaunchKernelEx( \ + &config, \ + per_token_group_quant_8bit_kernel, \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, num_groups, \ + groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \ + scale_num_rows, scale_stride); \ + } while (0) +#else + #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ + do { \ + per_token_group_quant_8bit_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, \ + num_groups, groups_per_block, (float)eps, (float)min_8bit, \ + (float)max_8bit, scale_num_rows, scale_stride); \ + } while (0) +#endif + +#define LAUNCH_KERNEL(T, DST_DTYPE) \ + do { \ + size_t smem_bytes = \ + static_cast(groups_per_block) * group_size * sizeof(T); \ + if (is_column_major) { \ + if (scale_ue8m0) { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \ + } \ + } else { \ + if (scale_ue8m0) { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \ + } \ + } \ } while (0) VLLM_STABLE_DISPATCH_FLOATING_TYPES( @@ -262,6 +283,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, })); #undef LAUNCH_KERNEL +#undef LAUNCH_KERNEL_INST } // Register-resident fast path for group_size==128. @@ -301,12 +323,21 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( const int sf_k_local = local_group_id % kGroupsPerBlockX; const int row_local = local_group_id / kGroupsPerBlockX; - const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local; - const int mn_idx = blockIdx.y * kRowsPerBlock + row_local; + // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap. + const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local; + const int mn_idx = blockIdx.x * kRowsPerBlock + row_local; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif if (mn_idx >= tma_aligned_mn) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif return; } + const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row); // Load 16 input elements (32 B) into registers as two adjacent uint4 @@ -417,6 +448,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( static_cast(mn_idx) * groups_per_row * GROUP_SIZE + sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE; *reinterpret_cast(group_output) = packed_out; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif } // Public entry point: register-resident packed quant kernel. @@ -473,6 +508,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, "]; got [", output_s_packed.stride(0), ", ", output_s_packed.stride(1), "]."); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); cudaStream_t stream = get_current_cuda_stream(); constexpr int THREADS_PER_GROUP = 8; @@ -484,34 +521,66 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, " is not a multiple of 4."); const int kx = GetGroupsPerBlockX(padded_groups_per_row); const int ry = 16 / kx; - const int64_t blocks_x = padded_groups_per_row / kx; - const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry; + const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry; + const int64_t sf_k_blocks = padded_groups_per_row / kx; const int num_threads = (kx * ry) * THREADS_PER_GROUP; - // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs. - STD_TORCH_CHECK(blocks_x <= static_cast(INT32_MAX) && - blocks_y <= static_cast(INT32_MAX), + // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535). + constexpr int64_t kMaxGridDimYZ = 65535; + STD_TORCH_CHECK(row_blocks <= static_cast(INT32_MAX) && + sf_k_blocks <= kMaxGridDimYZ, "per_token_group_quant_8bit_packed grid too large: (", - blocks_x, ", ", blocks_y, ")."); + row_blocks, ", ", sf_k_blocks, ")."); auto dst_type = output_q.scalar_type(); -#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \ - do { \ - dim3 grid(static_cast(blocks_x), \ - static_cast(blocks_y)); \ - dim3 block(num_threads); \ - per_token_group_quant_8bit_packed_register_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - reinterpret_cast(output_s_packed.data_ptr()), \ - static_cast(padded_groups_per_row), \ - static_cast(groups_per_row), static_cast(mn), \ - static_cast(output_q_mn_extent), \ - static_cast(tma_aligned_mn), num_scale_elems, \ - static_cast(eps), static_cast(min_8bit), \ - static_cast(max_8bit)); \ - } while (0) +// PDL (Programmatic Dependent Launch) is NVIDIA-only; ROCm/HIP has no +// equivalent launch attribute, so fall back to a classic launch there. +#ifndef USE_ROCM + #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \ + do { \ + cudaLaunchConfig_t config = {}; \ + config.gridDim = dim3(static_cast(row_blocks), \ + static_cast(sf_k_blocks)); \ + config.blockDim = dim3(num_threads); \ + 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, \ + per_token_group_quant_8bit_packed_register_kernel, \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(output_q_mn_extent), \ + static_cast(tma_aligned_mn), num_scale_elems, \ + static_cast(eps), static_cast(min_8bit), \ + static_cast(max_8bit)); \ + } while (0) +#else + #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \ + do { \ + dim3 grid(static_cast(row_blocks), \ + static_cast(sf_k_blocks)); \ + dim3 block(num_threads); \ + per_token_group_quant_8bit_packed_register_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(output_q_mn_extent), \ + static_cast(tma_aligned_mn), num_scale_elems, \ + static_cast(eps), static_cast(min_8bit), \ + static_cast(max_8bit)); \ + } while (0) +#endif #define LAUNCH_REG_KERNEL(T, DST_DTYPE) \ do { \ diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu b/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu index ede7913a355..e620c181847 100644 --- a/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu @@ -5,7 +5,7 @@ #include "../../../dispatch_utils.h" #include "../../../torch_utils.h" #include "../../vectorization_utils.cuh" -#include "../../../../cub_helpers.h" +#include "../../../cub_helpers.h" static inline __device__ int8_t float_to_int8_rn(float x) { #ifdef USE_ROCM diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index 68848b84566..519e213281c 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -665,6 +665,8 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, constexpr int kSortingAlgorithmThreshold = 12288; constexpr int kSplitWorkThreshold = 200 * 1000; constexpr int kNumThreadsPerBlock = 512; + const torch::stable::accelerator::DeviceGuard device_guard( + logits.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); const auto numColumns = logits.size(1); @@ -727,6 +729,8 @@ void top_k_per_row_prefill(const torch::stable::Tensor& logits, int64_t stride0, int64_t stride1, int64_t topK) { constexpr int kSortingAlgorithmThreshold = 12288; constexpr int kNumThreadsPerBlock = 512; + const torch::stable::accelerator::DeviceGuard device_guard( + logits.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); int numInsertionBlocks = diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 7656ba8cf8f..07b7521f863 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -21,6 +21,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, int64_t max_seq_len) { namespace P = vllm::persistent; + const torch::stable::accelerator::DeviceGuard device_guard( + logits.get_device_index()); const int64_t num_rows = logits.size(0); const int64_t stride = logits.stride(0); const cudaStream_t stream = get_current_cuda_stream(); @@ -260,6 +262,9 @@ void persistent_topk(const torch::stable::Tensor& logits, k == 512 || k == 1024 || k == 2048, "persistent_topk supports k=512, k=1024, or k=2048, got k=", k); + const torch::stable::accelerator::DeviceGuard device_guard( + logits.get_device_index()); + if (k == 512) { launch_persistent_topk<512>(logits, lengths, output, workspace, max_seq_len); diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh new file mode 100644 index 00000000000..5f9f823a339 --- /dev/null +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -0,0 +1,563 @@ +/* + * Shared 4096-bin single-CTA TopK helpers. + */ + +#ifndef TOPK_HISTOGRAM_4096_CUH_ +#define TOPK_HISTOGRAM_4096_CUH_ + +#include +#include +#include + +namespace vllm { +namespace topk_histogram_4096 { + +constexpr uint32_t kBlockSize = 1024; +constexpr uint32_t RADIX = 256; +constexpr uint32_t kMaxTies = 1024; +static_assert(kMaxTies <= kBlockSize, + "tie_handle requires kMaxTies <= kBlockSize"); +constexpr uint32_t kWarpSize = 32; +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; + +// Register path +constexpr uint32_t kHist4096VecsPerThread = 4; +constexpr uint32_t kHist4096MaxLen = + kHist4096VecsPerThread * 4 * kBlockSize; // 16384 + +struct alignas(16) MatchBin { + uint32_t bin, above_count, equal_count; +}; +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +__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; + const int p0 = (base < seq_len); + const int p1 = (base + 1 < seq_len); + const int p2 = (base + 2 < seq_len); + const 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); +} + +// converts the float32 score to a 32-bit ordered unsigned integer — the full +// precision key for radix sorting +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> +// bin 0-4095) +template +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { + __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 key >> (16 - kBits); +} + +// running sum within each warp — thread 0 gets its own value, thread 1 gets +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, + uint32_t v) { +#pragma unroll + for (uint32_t o = 1; o < 32; o *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); + if (lane >= o) v += n; + } + return v; +} + +// Returns the sum of a value across all 32 threads in the warp, and every +// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX +// instruction for hardware warp-wide reduction. Older targets use the +// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32 +// lanes). +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + uint32_t r; + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); + return r; +#else + #pragma unroll + for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { + v += __shfl_xor_sync(0xFFFFFFFF, v, mask); + } + return v; +#endif +} + +// ============================================================================ +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered +// key Each round narrows by 8 bits until ties are fully resolved +// ============================================================================ + +template +__device__ void tie_handle(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, void* _smem) { + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + + // Each thread loads one tie element. + const bool has = tx < num_ties; + const auto tie = has ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = convert_to_uint32_v2(tie.score); + + bool active = has; // tracks whether this thread's tie is still a candidate. + uint32_t remain = + TopK - num_above; // decreases each round as ties are resolved. + uint32_t wpos = TopK; // wpos will hold the final output position. + s->counter = 0; + __syncthreads(); + + // The 4-round radix loop - each round narrows by 8 bits until ties are fully + // resolved +#pragma unroll + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round + + // Step 1: Build 256-bin histogram. + if (tx < RADIX) s->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&s->histogram[bin], 1); + __syncthreads(); + + // Step 2: Prefix scan to find threshold + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; + } + __syncthreads(); + + if (tx < RADIX) { + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto tot = warp_reduce_sum_full(tmp); + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); + auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = {tx, above, remain - above}; + } + } + __syncthreads(); + + // Step 3: Scatter + auto [thr, na, _] = s->match; // threshold bin, num above, unused + if (active) { + if (bin > thr) { + wpos = num_above + + atomicAdd(&s->counter, 1); // above -> place in output directly + active = false; + } else if (bin < thr) + active = false; // below -> discard + else if (r == 3) + wpos = TopK - atomicAdd(&s->match.equal_count, + -1u); // last round: place remaining + } + remain -= na; + if (!remain) break; // all ties resolved early + } + // Final write + if (wpos < TopK) output[wpos] = tie.idx; +} + +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). +// tie_handle assumes 1 tie per thread (max 1024). +// This version handles 2 ties per thread via kPerThread=2 +template +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, + void* _smem) { + static_assert(TopK > kBlockSize); + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize; + const auto wi = tx / kWarpSize; + + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; + Tie my_ties[kPerThread]; + uint32_t keys[kPerThread]; + bool active[kPerThread]; + + for (uint32_t e = 0; e < kPerThread; e++) { + uint32_t idx = e * kBlockSize + tx; + if (idx < num_ties) { + my_ties[e] = ties[idx]; + keys[e] = convert_to_uint32_v2(ties[idx].score); + active[e] = true; + } else { + my_ties[e] = {0, 0.0f}; + keys[e] = 0; + active[e] = false; + } + } + + uint32_t remain = TopK - num_above; + s->counter = 0; + __syncthreads(); + + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; + if (tx < RADIX) { + s->histogram[tx] = 0; + } + __syncthreads(); + + for (uint32_t e = 0; e < kPerThread; e++) { + if (active[e]) { + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); + } + } + __syncthreads(); + + uint32_t hv = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + auto wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) { + s->warp_sum[wi] = wi2; + } + } + __syncthreads(); + if (tx < RADIX) { + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto total = warp_reduce_sum_full(tmp2); + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); + auto wi2 = warp_inclusive_sum(li, hv); + auto above = total - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = { + .bin = tx, .above_count = above, .equal_count = remain - above}; + } + } + __syncthreads(); + + auto thr = s->match.bin; + auto na = s->match.above_count; + + for (uint32_t e = 0; e < kPerThread; e++) { + if (!active[e]) { + continue; + } + uint32_t bin = (keys[e] >> sh) & 0xFF; + if (bin > thr) { + uint32_t wpos = num_above + atomicAdd(&s->counter, 1); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + active[e] = false; + } else if (bin < thr) { + active[e] = false; + } else if (r == 3) { + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + } + } + + num_above += na; + remain -= na; + __syncthreads(); + s->counter = 0; + __syncthreads(); + } +} + +// ============================================================================ +// Register-based single-CTA fast path for seq_len <= 16384 +// 4 float4 per thread × 1024 threads = 16384 elements max +// Uses 4096-bin (12-bit) histogram for better precision +// ============================================================================ + +template +struct Histogram4096Smem { + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[HIST_BINS]; + Tie tie_buffer[TIE_CAPACITY]; + }; +}; + +template +__device__ void histogram_4096_topk(const float* __restrict__ scores, + int32_t* __restrict__ output, + uint32_t length, void* _smem) { + constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; + static_assert(HIST_BINS >= kBlockSize, + "HIST_BITS must give >= kBlockSize bins"); + + using Smem = Histogram4096Smem; + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + + // Phase 1: Load all data into RF + build histogram + float4 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread + if constexpr (ITEMS_PER_THREAD >= 4) { + // Zero the histogram (SMEM writes) + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) + reinterpret_cast( + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = + make_uint4(0, 0, 0, 0); + } else { + if (tx < HIST_BINS) smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + if constexpr (UsePredicatedLoads) { + const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + if (row_aligned && base + 3 < length) { + vecs[v] = *reinterpret_cast(scores + base); + } else { + load_float4_predicated(scores + base, static_cast(base), + static_cast(length), vecs[v].x, vecs[v].y, + vecs[v].z, vecs[v].w); + } + } + } + } else { +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + vecs[v] = *reinterpret_cast(scores + base); + } + } + } + __syncthreads(); + + // Build histogram from RF via atomic adds into the shared histogram + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], + 1); + } + } + } + __syncthreads(); + + // Phase 2: Prefix scan to find threshold bin + // Multi-element scan (4096 bins: 4 per thread) + uint32_t orig[ITEMS_PER_THREAD]; + uint32_t local_sum = 0; + + // Step 1: Each thread sums its 4 bins +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; + local_sum += orig[i]; + } + + // Step 2: Warp-level inclusive prefix sum on local_sum + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + // Step 3: Inter-warp prefix across warp sums. + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix = warp_reduce_sum_full( + lane_id < warp_id ? tmp : 0); // sum of all prior warps + prefix += + warp_inc - local_sum; // exclusive prefix within this thread's position + + // Step 4: Find threshold - scan 4 bins, accumulate prefix +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + prefix += orig[i]; + const auto above = length - prefix; // elements in bins ABOVE this one + if (above < TopK && above + orig[i] >= TopK) { + smem->match = {.bin = tx * ITEMS_PER_THREAD + i, + .above_count = above, + .equal_count = orig[i]}; + } + } + + __syncthreads(); + + // Phase 3: Scatter from registers + const auto [thr_bin, num_above, num_equal] = smem->match; + const bool need_tie = (num_equal + num_above > TopK); + + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + + // Fast warp-ballot tie-breaking for small tie counts + const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t topk_remain = + TopK - num_above; // pick exactly remaining elements to fill topK + + auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpSize) { + // <=32 ties - Use warp ballot + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 + // comparisons in one instruction per warp. O(1) work. + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie + const auto target = + smem->tie_buffer[warp_id]; // each warp evaluates one candidate + const bool pred = + is_greater(tie, target); // compare all ties against target + const auto rank = static_cast( + __popc(__ballot_sync(mask, pred))); // count how many are greater + if (lane_id == 0 && rank < topk_remain) { + output[num_above + rank] = target.idx; // place at correct position + } + } else if (num_ties <= + kWarpSize * + 2) { // TODO (roberto): try to refactor this with <=32 case + // Same idea but each thread handles 2 tie elements + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + const auto lane1 = lane_id + kWarpSize; + const auto warp1 = warp_id + kWarpSize; + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; + const auto tie0 = smem->tie_buffer[lane_id]; + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; + if (warp_id < num_ties) { + const auto target = smem->tie_buffer[warp_id]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + if (warp1 < num_ties) { + const auto target = smem->tie_buffer[warp1]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + } else { + // Large tie count: fall back to 4-round radix-256 sort + if constexpr (TopK <= kBlockSize) { + tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + } else { + tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + smem); + } + } +} + +template +__device__ __noinline__ void histogram_4096_topk_predicated( + const float* __restrict__ scores, int32_t* __restrict__ output, + uint32_t length, void* _smem) { + histogram_4096_topk(scores, output, + length, _smem); +} + +} // namespace topk_histogram_4096 +} // namespace vllm + +#endif // TOPK_HISTOGRAM_4096_CUH_ diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index e7d1b3669fb..fd2ffe92e08 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -29,10 +29,85 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "()"); ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + #ifndef USE_ROCM - // TODO: Remove this once ROCm upgrade to torch 2.11. - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + // Note about marlin kernel 'workspace' arguments: + // Technically these should be mutable since they are modified by the kernel. + // But since they are set back to zero once the kernel is finished we can + // hand wave and say that they have no net effect. + // + // The reason to mark 'workspace' as immutable is so that they don't interfere + // with using ScalarType arguments in the ops. If they are marked as mutable, + // pytorch throws an assert in + // 'torch._higher_order_ops._register_effectful_op' that prevents these + // kernels from being torch.compile'd. + // See the following document for more info on custom types and ops that use + // custom types: + // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA + + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. + ops.def( + "machete_supported_schedules(" + " ScalarType a_type," + " int b_type," + " ScalarType? maybe_group_scales_type," + " ScalarType? maybe_group_zeros_type," + " ScalarType? maybe_channel_scales_type," + " ScalarType? maybe_token_scales_type," + " ScalarType? maybe_out_type" + ") -> str[]"); + ops.def( + "machete_mm(" + " Tensor A," + " Tensor B," + " int b_type," + " ScalarType? out_type," + " Tensor? group_scales," + " Tensor? group_zeros," + " int? group_size," + " Tensor? channel_scales," + " Tensor? token_scales," + " str? schedule" + ") -> Tensor"); + ops.def( + "machete_prepack_B(" + " Tensor B," + " ScalarType a_type," + " int b_type," + " ScalarType? group_scales_type" + ") -> Tensor"); + // conditionally compiled so impl registration is in source file + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor? " + "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file #endif #ifndef USE_ROCM @@ -291,12 +366,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // Apply Root Mean Square (RMS) Normalization to the input tensor. ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> " + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " "()"); // In-place fused Add and RMS Normalization. ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, " + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " "float epsilon) -> ()"); // Layernorm-quant @@ -372,10 +448,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int cache_block_size) -> ()"); #ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms(" - "Tensor input, Tensor norm_weight, Tensor workspace, " - "int rank, int nranks, float eps) -> Tensor"); ops.def( "minimax_allreduce_rms_qk(" "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " @@ -383,6 +455,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "float eps) -> (Tensor, Tensor)"); #endif + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. + ops.def( + "fused_minimax_m3_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor cos_sin_cache, Tensor positions, int num_heads, " + "int num_kv_heads, int rotary_dim, float eps, " + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " + "int num_index_heads, " + "Tensor? slot_mapping, Tensor? index_slot_mapping, " + "Tensor!? kv_cache, Tensor!? index_cache, " + "int block_size, Tensor!? q_out, Tensor!? index_q_out, " + "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); + // Apply repetition penalties to logits in-place. ops.def( "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " @@ -403,16 +488,33 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " "Tensor workspace, int k, int max_seq_len) -> ()"); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + // Activation ops + ops.def( + "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " + "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); + ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + // Activation function used in SwiGLU. ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); // SwiGLU activation with input clamping. + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. ops.def( - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) " - "-> ()"); + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " + "float alpha=1.0, float beta=0.0) -> ()"); + + // SwiGLU activation with FP8 quantization. + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); // Activation function used in GeGLU with `none` approximation. ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); @@ -437,6 +539,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // Quick GELU implementation. ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); + // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 + ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); + // Compute int8 quantized tensor for given scaling factor. ops.def( "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," @@ -479,34 +584,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // Post processing for GPTQ. ops.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); - // Dequantization for GGML. - ops.def( - "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? " - "dtype) -> Tensor"); - - // mmvq kernel for GGML. - ops.def( - "ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) " - "-> Tensor"); - - // mmq kernel for GGML. - ops.def( - "ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor"); - - // moe kernel for GGML. - ops.def( - "ggml_moe_a8(Tensor X, Tensor W, " - "Tensor sorted_token_ids, Tensor expert_ids, Tensor " - "num_tokens_post_padded, " - "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor"); - - ops.def( - "ggml_moe_a8_vec(Tensor X, Tensor W, " - "Tensor topk_ids, int top_k, " - "int type, SymInt row, SymInt tokens) -> Tensor"); - - ops.def("ggml_moe_get_block_size(int type) -> int"); - // Mamba selective scan kernel ops.def( "selective_scan_fwd(Tensor! u, Tensor! delta," @@ -525,35 +602,21 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? cu_chunk_seqlen," "Tensor? last_chunk_indices) -> ()"); - // Attention ops - // Compute the attention between an input query and the cached - // keys/values using PagedAttention. + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. ops.def( - "paged_attention_v1(" - " Tensor! out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - - // PagedAttention V2. - ops.def( - "paged_attention_v2(" - " Tensor! out, Tensor! exp_sums, Tensor! max_logits," - " Tensor! tmp_out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + // Per-token group quantization ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); ops.impl("per_token_group_fp8_quant_packed", @@ -626,9 +689,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); #ifndef USE_ROCM - ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms)); ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif + ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", + TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); // Sampler kernels (shared CUDA/ROCm) ops.impl("apply_repetition_penalties_", @@ -636,8 +700,15 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif // Activation kernels (shared CUDA/ROCm) + ops.impl("persistent_masked_m_silu_mul_quant", + TORCH_BOX(&persistent_masked_m_silu_mul_quant)); + ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); + ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); @@ -647,6 +718,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); + ops.impl("relu_squared", TORCH_BOX(&relu_squared)); ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); // INT8 quantization kernels @@ -663,20 +735,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); - // GGML kernels - ops.impl("ggml_dequantize", TORCH_BOX(&ggml_dequantize)); - ops.impl("ggml_mul_mat_vec_a8", TORCH_BOX(&ggml_mul_mat_vec_a8)); - ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8)); - ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); - ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); + // Mamba kernels ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); - - ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1)); - ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2)); } -// TODO: Remove this once ROCm upgrade to torch 2.11. -#ifndef USE_ROCM STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { ops.impl("get_cuda_view_from_cpu_tensor", TORCH_BOX(&get_cuda_view_from_cpu_tensor)); @@ -695,8 +757,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); } -#endif - // These capability-check functions take only primitive args (no tensors), so // there is no device to dispatch on. CompositeExplicitAutograd makes them // available for all backends. This is the stable ABI equivalent of calling @@ -712,9 +772,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { ops.impl("cutlass_scaled_mm_supports_fp4", TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); #endif - - // GGML block size lookup (no tensor args) - ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size)); } // Cache ops diff --git a/csrc/type_convert.cuh b/csrc/libtorch_stable/type_convert.cuh similarity index 100% rename from csrc/type_convert.cuh rename to csrc/libtorch_stable/type_convert.cuh diff --git a/csrc/moe/dsv3_router_gemm_entry.cu b/csrc/moe/dsv3_router_gemm_entry.cu deleted file mode 100644 index 38fb681c223..00000000000 --- a/csrc/moe/dsv3_router_gemm_entry.cu +++ /dev/null @@ -1,169 +0,0 @@ -/* - * Adapted from SGLang's sgl-kernel implementation, which was adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, 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 "core/registration.h" -#include "dsv3_router_gemm_utils.h" - -static constexpr int DEFAULT_NUM_EXPERTS = 256; -static constexpr int KIMI_K2_NUM_EXPERTS = 384; -static constexpr int DEFAULT_HIDDEN_DIM = 7168; - -template -void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, - cudaStream_t stream); - -template -void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, - T const* mat_b, cudaStream_t stream); - -template -struct LoopUnroller { - static void unroll_float_output(int num_tokens, float* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts, - kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_float_output(num_tokens, output, input, - weights, stream); - } - } - - static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts, - kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_bf16_output(num_tokens, output, input, - weights, stream); - } - } -}; - -template -struct LoopUnroller { - static void unroll_float_output(int num_tokens, float* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( - output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } - - static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( - output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } -}; - -void dsv3_router_gemm(at::Tensor& output, // [num_tokens, num_experts] - const at::Tensor& mat_a, // [num_tokens, hidden_dim] - const at::Tensor& mat_b // [num_experts, hidden_dim] -) { - TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); - - const int num_tokens = mat_a.size(0); - const int num_experts = mat_b.size(0); - const int hidden_dim = mat_a.size(1); - - TORCH_CHECK(mat_a.size(1) == mat_b.size(1), - "mat_a and mat_b must have the same hidden_dim"); - TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM, - "Expected hidden_dim=", DEFAULT_HIDDEN_DIM, - ", but got hidden_dim=", hidden_dim); - TORCH_CHECK( - num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS, - "Expected num_experts=", DEFAULT_NUM_EXPERTS, - " or num_experts=", KIMI_K2_NUM_EXPERTS, - ", but got num_experts=", num_experts); - TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, - "currently num_tokens must be less than or equal to 16 for " - "router_gemm"); - TORCH_CHECK(mat_a.dtype() == at::kBFloat16, "mat_a must be bf16"); - TORCH_CHECK(mat_b.dtype() == at::kBFloat16, "mat_b must be bf16"); - TORCH_CHECK(output.dtype() == at::kFloat || output.dtype() == at::kBFloat16, - "output must be float32 or bf16"); - - auto const sm = getSMVersion(); - TORCH_CHECK(sm >= 90 && sm <= 103, "required SM_103 >= CUDA ARCH >= SM_90"); - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - if (output.dtype() == at::kFloat) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>:: - unroll_float_output( - num_tokens, reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); - } else if (num_experts == KIMI_K2_NUM_EXPERTS) { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>:: - unroll_float_output( - num_tokens, reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); - } - } else if (output.dtype() == at::kBFloat16) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>:: - unroll_bf16_output( - num_tokens, - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); - } else if (num_experts == KIMI_K2_NUM_EXPERTS) { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>:: - unroll_bf16_output( - num_tokens, - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); - } - } -} - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("dsv3_router_gemm", &dsv3_router_gemm); -} diff --git a/csrc/moe/dsv3_router_gemm_utils.h b/csrc/moe/dsv3_router_gemm_utils.h deleted file mode 100644 index 9b533bcabfc..00000000000 --- a/csrc/moe/dsv3_router_gemm_utils.h +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Adapted from SGLang's sgl-kernel implementation, which was adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, 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 - -inline int getSMVersion() { - auto* props = at::cuda::getCurrentDeviceProperties(); - return props->major * 10 + props->minor; -} diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp index 58dc4020168..1b071d334ff 100644 --- a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp +++ b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp @@ -29,25 +29,37 @@ enum ActivationKind : int64_t { torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind) { + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind) { TORCH_CHECK(x.dim() == 2, "x must be 2D"); TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2, "topk tensors must be [T, K]"); TORCH_CHECK( w13_packed.size(0) == w2_packed.size(0), "w13_packed and w2_packed must have same number of experts in dim 0"); - TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I"); const int64_t T = x.size(0); const int64_t K = topk_ids.size(1); const int64_t E = w13_packed.size(0); const int64_t N = T * K; + const int64_t w13_out_features = 2 * intermediate_size; auto x_c = x.contiguous(); + // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation + // dtypes: + // - fp32: with channelwise and groupwise + // - bf16: with channelwise -> upcast to fp32 for groupwise + // - fp16: not supported -> upcast to fp32 for groupwise & channelwise + const auto output_dtype = x_c.scalar_type(); + const bool should_cast_input = + ((group_size != -1) && output_dtype == at::kBFloat16) || + output_dtype == at::kHalf; + if (should_cast_input) { + x_c = x_c.to(at::kFloat); + } auto ids_c = topk_ids.contiguous(); - auto gates_c = topk_weights.to(at::kFloat).contiguous(); + auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous(); // bucketing tokens -> experts c10::SmallVector counts( @@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu( c10::SmallVector offsets(E + 1, 0); // ( E +1 ) for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e]; + // expert_tokens = [tokens indices for expert 0, ...] + // expert_gates = [router weights for tokens assigned to expert 0, ...] auto expert_tokens = at::empty({offsets[E]}, ids_c.options()); auto expert_gates = at::empty({offsets[E]}, gates_c.options()); { c10::SmallVector cursor(E, 0); - const auto* ids_ptr = ids_c.data_ptr(); - const auto* gts_ptr = gates_c.data_ptr(); - auto* tok_ptr = expert_tokens.data_ptr(); - auto* gate_ptr = expert_gates.data_ptr(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(), + "bucket_expert_tokens_and_gates", [&] { + const auto* ids_ptr = ids_c.data_ptr(); + const auto* gts_ptr = gates_c.data_ptr(); + auto* tok_ptr = expert_tokens.data_ptr(); + auto* gate_ptr = expert_gates.data_ptr(); - for (int64_t t = 0; t < T; ++t) { - const int64_t base = t * K; - for (int64_t k = 0; k < K; ++k) { - const int64_t idx = base + k; - const int64_t e = ids_ptr[idx]; - const int64_t p = offsets[e] + (cursor[e]++); - tok_ptr[p] = t; - gate_ptr[p] = gts_ptr[idx]; - } - } + for (int64_t t = 0; t < T; ++t) { + const int64_t base = t * K; + for (int64_t k = 0; k < K; ++k) { + const int64_t idx = base + k; + const int64_t e = ids_ptr[idx]; + const int64_t p = offsets[e] + (cursor[e]++); + tok_ptr[p] = t; + gate_ptr[p] = gts_ptr[idx]; + } + } + }); } - const int64_t g_eff_13 = (group_size != -1) ? group_size : H; - const int64_t g_eff_2 = (group_size != -1) ? group_size : I; + const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size; + const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size; + // X_all [num_tokens * K, hidden_size] auto X_all = x_c.index_select(/*dim=*/0, expert_tokens); if (apply_router_weight_on_input) { X_all = X_all.mul(expert_gates.unsqueeze(1)); } - auto Y_all = at::empty({offsets[E], H}, x_c.options()); + auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options()); at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) { c10::InferenceMode guard; @@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu( auto w2_e = w2_packed.select(/*dim=*/0, e); // W13 - auto y13 = - mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2); + auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size, + /*out_features=*/w13_out_features); - auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I); - auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I); + auto g_part = + y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size); + auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size, + /*length=*/intermediate_size); torch::Tensor act; if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI @@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu( } // W2 - auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H); + auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size, + /*out_features=*/hidden_size); // Store per-expert result Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y); @@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu( if (!apply_router_weight_on_input) { Y_all = Y_all.mul(expert_gates.unsqueeze(1)); } + if (Y_all.scalar_type() != output_dtype) { + Y_all = Y_all.to(output_dtype); + } - auto out = at::zeros({T, H}, x.options()); + auto out = at::zeros({T, hidden_size}, x.options()); out = at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all); diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h deleted file mode 100644 index ca2776c6edd..00000000000 --- a/csrc/moe/moe_ops.h +++ /dev/null @@ -1,81 +0,0 @@ -#pragma once - -#include - -void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices, - torch::Tensor& token_expert_indices, - torch::Tensor& gating_output, bool renormalize, - std::optional bias); - -void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices, - torch::Tensor& token_expert_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, - int64_t block_size, torch::Tensor sorted_token_ids, - torch::Tensor experts_ids, - torch::Tensor num_tokens_post_pad, - std::optional maybe_expert_map); - -void batched_moe_align_block_size(int64_t max_tokens_per_batch, - int64_t block_size, - torch::Tensor const& expert_num_tokens, - torch::Tensor sorted_ids, - torch::Tensor expert_ids, - torch::Tensor num_tokens_post_pad); - -void moe_lora_align_block_size( - torch::Tensor topk_ids, torch::Tensor token_lora_mapping, - int64_t num_experts, int64_t block_size, int64_t max_loras, - int64_t max_num_tokens_padded, int64_t max_num_m_blocks, - torch::Tensor sorted_token_ids, torch::Tensor expert_ids, - torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled, - torch::Tensor lora_ids, std::optional maybe_expert_map); -#ifndef USE_ROCM -torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output, - torch::Tensor b_qweight, torch::Tensor b_scales, - std::optional b_qzeros, - std::optional topk_weights, - torch::Tensor sorted_token_ids, - torch::Tensor expert_ids, - torch::Tensor num_tokens_post_pad, int64_t top_k, - int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, - int64_t BLOCK_SIZE_K, int64_t bit); - -std::tuple grouped_topk( - torch::Tensor const& scores, int64_t n_group, int64_t topk_group, - int64_t topk, bool renormalize, double routed_scaling_factor, - torch::Tensor const& bias, int64_t scoring_func); -#endif - -bool moe_permute_unpermute_supported(); - -int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, - int64_t num_experts); - -void shuffle_rows(const torch::Tensor& input_tensor, - const torch::Tensor& dst2src_map, - torch::Tensor& output_tensor); - -#ifndef USE_ROCM -// DeepSeek V3 optimized router GEMM kernel for SM90+ -// Computes output = mat_a @ mat_b.T where: -// mat_a: [num_tokens, hidden_dim] in bf16 -// mat_b: [num_experts, hidden_dim] in bf16 -// output: [num_tokens, num_experts] in bf16 or fp32 -// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168 -void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, - const torch::Tensor& mat_b); -#endif diff --git a/csrc/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu deleted file mode 100644 index 6fce009ae6d..00000000000 --- a/csrc/moe/moe_permute_unpermute_op.cu +++ /dev/null @@ -1,286 +0,0 @@ -#include -#include -#include -#include "permute_unpermute_kernels/moe_permute_unpermute_kernel.h" -#include "permute_unpermute_kernels/dispatch.h" -#include "core/registration.h" - -// moe_permute kernels require at least CUDA 12.0 -#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) - -namespace { - -torch::Tensor maybe_allocate_tensor( - const std::optional& maybe_tensor, - at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device, - char const* name) { - auto expected_numel = c10::multiply_integers(expected_sizes); - if (maybe_tensor.has_value()) { - auto tensor = maybe_tensor.value(); - TORCH_CHECK(tensor.device() == device, name, " must be on the same device"); - TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype"); - TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); - TORCH_CHECK(tensor.numel() >= expected_numel, name, - " is too small for the requested shape"); - auto flat_tensor = tensor.view({tensor.numel()}); - return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes); - } - return torch::empty(expected_sizes, torch::dtype(dtype).device(device)); -} - -} // namespace - -int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, - int64_t n_expert) { - return static_cast( - CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert)); -} - -void moe_permute_impl( - const torch::Tensor& input, // [n_token, hidden] - const torch::Tensor& topk_ids, // [n_token, topk] - const torch::Tensor& token_expert_indices, // [n_token, topk] - const std::optional& expert_map, // [n_expert] - int64_t n_expert, int64_t n_local_expert, int64_t topk, - torch::Tensor& permuted_input, // [permuted_size, hidden] - torch::Tensor& expert_first_token_offset, // [n_local_expert + 1] - torch::Tensor& inv_permuted_idx, // [n_token, topk] - torch::Tensor& permuted_idx, // [permute_size] - const std::optional& maybe_sort_workspace, - const std::optional& maybe_permuted_experts_id, - const std::optional& maybe_sorted_row_idx, - const std::optional& maybe_topk_ids_for_sort) { - TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long, - "expert_first_token_offset must be int64"); - TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int, - "topk_ids must be int32"); - TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int, - "token_expert_indices must be int32"); - TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int, - "inv_permuted_idx must be int32"); - TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1, - "expert_first_token_offset shape != n_local_expert+1"); - TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(), - "token_expert_indices shape must be same as inv_permuted_idx"); - auto device = input.device(); - auto n_token = input.sizes()[0]; - auto n_hidden = input.sizes()[1]; - auto expanded_rows = n_token * topk; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert); - auto sort_workspace = - maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8, - device, "sort_workspace"); - auto permuted_experts_id = - maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(), - at::ScalarType::Int, device, "permuted_experts_id"); - auto sorted_row_idx = - maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(), - at::ScalarType::Int, device, "sorted_row_idx"); - - CubKeyValueSorter sorter{}; - int64_t* valid_num_ptr = nullptr; - torch::Tensor topk_ids_for_sort = topk_ids; - - if (expert_map.has_value()) { - const int* expert_map_ptr = get_ptr(expert_map.value()); - valid_num_ptr = - get_ptr(expert_first_token_offset) + n_local_expert; - topk_ids_for_sort = - maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(), - at::ScalarType::Int, device, "topk_ids_for_sort"); - topk_ids_for_sort.copy_(topk_ids); - preprocessTopkIdLauncher(get_ptr(topk_ids_for_sort), n_token * topk, - expert_map_ptr, n_expert, stream); - } - - sortAndScanExpert( - get_ptr(topk_ids_for_sort), get_ptr(token_expert_indices), - get_ptr(permuted_experts_id), get_ptr(sorted_row_idx), - get_ptr(expert_first_token_offset), n_token, n_expert, - n_local_expert, topk, sorter, get_ptr(sort_workspace), stream); - - MOE_DISPATCH(input.scalar_type(), [&] { - expandInputRowsKernelLauncher( - get_ptr(input), get_ptr(permuted_input), - get_ptr(sorted_row_idx), get_ptr(inv_permuted_idx), - get_ptr(permuted_idx), get_ptr(expert_first_token_offset), - n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream); - }); -} - -void moe_permute( - const torch::Tensor& input, // [n_token, hidden] - const torch::Tensor& topk_ids, // [n_token, topk] - const torch::Tensor& token_expert_indices, // [n_token, topk] - const std::optional& expert_map, // [n_expert] - int64_t n_expert, int64_t n_local_expert, int64_t topk, - torch::Tensor& permuted_input, // [permuted_size, hidden] - torch::Tensor& expert_first_token_offset, // [n_local_expert + 1] - torch::Tensor& inv_permuted_idx, // [n_token, topk] - torch::Tensor& permuted_idx) { // [permute_size] - moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, - n_local_expert, topk, permuted_input, - expert_first_token_offset, inv_permuted_idx, permuted_idx, - std::nullopt, std::nullopt, std::nullopt, std::nullopt); -} - -void moe_permute_with_scratch( - const torch::Tensor& input, const torch::Tensor& topk_ids, - const torch::Tensor& token_expert_indices, - const std::optional& expert_map, int64_t n_expert, - int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input, - torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx, - torch::Tensor& permuted_idx, torch::Tensor& sort_workspace, - torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx, - torch::Tensor& topk_ids_for_sort) { - moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, - n_local_expert, topk, permuted_input, - expert_first_token_offset, inv_permuted_idx, permuted_idx, - sort_workspace, permuted_experts_id, sorted_row_idx, - topk_ids_for_sort); -} - -void moe_unpermute( - const torch::Tensor& permuted_hidden_states, // [n_token * topk, hidden] - const torch::Tensor& topk_weights, // [n_token, topk] - const torch::Tensor& inv_permuted_idx, // [n_token, topk] - const std::optional& - expert_first_token_offset, // [n_local_expert+1] - int64_t topk, - torch::Tensor& hidden_states // [n_token, hidden] -) { - TORCH_CHECK( - permuted_hidden_states.scalar_type() == hidden_states.scalar_type(), - "permuted_hidden_states dtype must be same as hidden_states"); - auto n_token = hidden_states.size(0); - auto n_hidden = hidden_states.size(1); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - int64_t const* valid_ptr = nullptr; - if (expert_first_token_offset.has_value()) { - int n_local_expert = expert_first_token_offset.value().size(0) - 1; - valid_ptr = - get_ptr(expert_first_token_offset.value()) + n_local_expert; - } - - MOE_DISPATCH(hidden_states.scalar_type(), [&] { - finalizeMoeRoutingKernelLauncher( - get_ptr(permuted_hidden_states), - get_ptr(hidden_states), get_ptr(topk_weights), - get_ptr(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr, - stream); - }); -} - -template -__global__ void shuffleInputRowsKernel(const T* input, - const int32_t* dst2src_map, T* output, - int64_t num_src_rows, - int64_t num_dst_rows, int64_t num_cols) { - int64_t dest_row_idx = blockIdx.x; - int64_t const source_row_idx = dst2src_map[dest_row_idx]; - - if (blockIdx.x < num_dst_rows) { - // Load 128-bits per thread - constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8; - using DataElem = cutlass::Array; - - // Duplicate and permute rows - auto const* source_row_ptr = - reinterpret_cast(input + source_row_idx * num_cols); - auto* dest_row_ptr = - reinterpret_cast(output + dest_row_idx * num_cols); - - int64_t const start_offset = threadIdx.x; - int64_t const stride = blockDim.x; - int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD; - - for (int elem_index = start_offset; elem_index < num_elems_in_col; - elem_index += stride) { - dest_row_ptr[elem_index] = source_row_ptr[elem_index]; - } - } -} - -void shuffle_rows(const torch::Tensor& input_tensor, - const torch::Tensor& dst2src_map, - torch::Tensor& output_tensor) { - TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(), - "Input and output tensors must have the same data type"); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - int64_t const blocks = output_tensor.size(0); - int64_t const threads = 256; - int64_t const num_dest_rows = output_tensor.size(0); - int64_t const num_src_rows = input_tensor.size(0); - int64_t const num_cols = input_tensor.size(1); - - TORCH_CHECK(!(num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)), - "num_cols must be divisible by 128 / " - "sizeof(input_tensor.scalar_type()) / 8"); - - MOE_DISPATCH(input_tensor.scalar_type(), [&] { - shuffleInputRowsKernel<<>>( - reinterpret_cast(input_tensor.data_ptr()), - dst2src_map.data_ptr(), - reinterpret_cast(output_tensor.data_ptr()), num_src_rows, - num_dest_rows, num_cols); - }); -} - -#else - -int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, - int64_t n_expert) { - TORCH_CHECK( - false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0"); -} - -void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids, - const torch::Tensor& token_expert_indices, - const std::optional& expert_map, - int64_t n_expert, int64_t n_local_expert, int64_t topk, - torch::Tensor& permuted_input, - torch::Tensor& expert_first_token_offset, - torch::Tensor& inv_permuted_idx, torch::Tensor& permuted_idx) { - TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0"); -} - -void moe_permute_with_scratch( - const torch::Tensor& input, const torch::Tensor& topk_ids, - const torch::Tensor& token_expert_indices, - const std::optional& expert_map, int64_t n_expert, - int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input, - torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx, - torch::Tensor& permuted_idx, torch::Tensor& sort_workspace, - torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx, - torch::Tensor& topk_ids_for_sort) { - TORCH_CHECK(false, - "moe_permute_with_scratch is not supported on CUDA < 12.0"); -} - -void moe_unpermute( - const torch::Tensor& permuted_hidden_states, - const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx, - const std::optional& expert_first_token_offset, int64_t topk, - torch::Tensor& hidden_states) { - TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0"); -} - -#endif - -bool moe_permute_unpermute_supported() { -#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) - return true; -#else - return false; -#endif -} - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("moe_permute", &moe_permute); - m.impl("moe_permute_with_scratch", &moe_permute_with_scratch); - m.impl("moe_unpermute", &moe_unpermute); -} \ No newline at end of file diff --git a/csrc/moe/permute_unpermute_kernels/dispatch.h b/csrc/moe/permute_unpermute_kernels/dispatch.h deleted file mode 100644 index d0f1ea4aded..00000000000 --- a/csrc/moe/permute_unpermute_kernels/dispatch.h +++ /dev/null @@ -1,59 +0,0 @@ -#pragma once -#include -#define MOE_SWITCH(TYPE, ...) \ - at::ScalarType _st = ::detail::scalar_type(TYPE); \ - switch (_st) { \ - __VA_ARGS__ \ - default: \ - TORCH_CHECK(false, "[moe permute]data type dispatch fail!") \ - } - -#define MOE_DISPATCH_CASE(enum_type, ...) \ - case enum_type: { \ - using scalar_t = ScalarType2CudaType::type; \ - __VA_ARGS__(); \ - break; \ - } -#define MOE_DISPATCH_FLOAT_CASE(...) \ - MOE_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - MOE_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - MOE_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - MOE_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \ - MOE_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - MOE_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) - -#define MOE_DISPATCH(TYPE, ...) \ - MOE_SWITCH(TYPE, MOE_DISPATCH_FLOAT_CASE(__VA_ARGS__)) - -template -struct ScalarType2CudaType; - -template <> -struct ScalarType2CudaType { - using type = float; -}; -template <> -struct ScalarType2CudaType { - using type = half; -}; -template <> -struct ScalarType2CudaType { - using type = __nv_bfloat16; -}; -// uint8 for packed fp4 -template <> -struct ScalarType2CudaType { - using type = uint8_t; -}; - -// #if __CUDA_ARCH__ >= 890 -// fp8 -template <> -struct ScalarType2CudaType { - using type = __nv_fp8_e5m2; -}; -template <> -struct ScalarType2CudaType { - using type = __nv_fp8_e4m3; -}; -// #endif \ No newline at end of file diff --git a/csrc/ops.h b/csrc/ops.h index e39bae08f19..d7ee2d080c4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -9,36 +9,14 @@ #include -torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { - // Ensure tensor is on CUDA - if (!tensor.is_cuda()) { - throw std::runtime_error("Tensor must be on CUDA device"); - } - - // Get the raw data pointer - void* data_ptr = tensor.data_ptr(); - - // Get tensor sizes and strides - std::vector sizes = tensor.sizes().vec(); - std::vector strides = tensor.strides().vec(); - - // Get tensor options (dtype, device) - auto options = tensor.options(); - - // Create a new tensor from the raw data pointer - auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options); - - return new_tensor; -} - // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. -void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, - double epsilon); +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon); void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - torch::Tensor& weight, double epsilon); + std::optional weight, double epsilon); // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable // ABI for CUDA). It remains here because the CPU build still uses these @@ -50,33 +28,22 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, void silu_and_mul(torch::Tensor& out, torch::Tensor& input); -void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit); - -void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input, - torch::Tensor& scale); - -void persistent_masked_m_silu_mul_quant( - const at::Tensor& input, // (E, T, 2*H) - const at::Tensor& counts, // (E) - at::Tensor& y_q, // (E, T, H) [OUT] - at::Tensor& y_s, // (E, T, H//group_size) [OUT] - bool use_ue8m0); +void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit, + double alpha = 1.0, double beta = 0.0); void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void gelu_tanh(torch::Tensor& out, torch::Tensor& input); + void gelu_new(torch::Tensor& out, torch::Tensor& input); void gelu_fast(torch::Tensor& out, torch::Tensor& input); void gelu_quick(torch::Tensor& out, torch::Tensor& input); -void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope, - torch::Tensor const& q_pe, - torch::Tensor const& kv_c_and_k_pe_cache, - torch::Tensor const& seq_lens, - torch::Tensor const& page_table, double scale); +void relu_squared(torch::Tensor& out, torch::Tensor& input); void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale, @@ -88,9 +55,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind); + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); using fptr_t = int64_t; #ifdef USE_ROCM @@ -102,7 +69,4 @@ void qr_open_handles(fptr_t _fa, const std::vector& handles); void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, int64_t quant_level, bool cast_bf2half = false); int64_t qr_max_size(); - -// TODO: Remove this once ROCm upgrade to torch 2.11. -torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor); #endif diff --git a/csrc/quantization/machete/machete_mm_launcher.cuh b/csrc/quantization/machete/machete_mm_launcher.cuh deleted file mode 100644 index cabe0af46f0..00000000000 --- a/csrc/quantization/machete/machete_mm_launcher.cuh +++ /dev/null @@ -1,75 +0,0 @@ -#pragma once - -#include -#include - -#include "machete_mm_kernel.cuh" -#include "cutlass_extensions/torch_utils.hpp" -#include "core/scalar_type.hpp" - -namespace machete { - -struct MMArgs { - torch::Tensor const& A; - torch::Tensor const& B; - vllm::ScalarType const& b_type; - std::optional const& maybe_out_type; - std::optional const& maybe_group_scales; - std::optional const& maybe_group_zeros; - std::optional maybe_group_size; - std::optional const& maybe_channel_scales; - std::optional const& maybe_token_scales; - std::optional maybe_schedule; -}; - -struct SupportedSchedulesArgs { - at::ScalarType a_type; - vllm::ScalarType b_type; - std::optional maybe_group_scales_type; - std::optional maybe_group_zeros_type; - std::optional maybe_channel_scales_type; - std::optional maybe_token_scales_type; - std::optional maybe_out_type; -}; - -torch::Tensor mm_dispatch(MMArgs args); - -std::vector supported_schedules_dispatch( - SupportedSchedulesArgs args); - -template -torch::Tensor run_impl(MMArgs args) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(args.A)); - - auto device = args.A.device(); - auto stream = at::cuda::getCurrentCUDAStream(device.index()); - - int M = args.A.size(0); - int N = args.B.size(1); - int K = args.A.size(1); - - // Allocate output - torch::Tensor D = torch::empty( - {M, N}, - torch::TensorOptions() - .dtype(equivalent_scalar_type_v) - .device(device)); - - auto arguments = MacheteKernel::create_arguments( - stream, // - args.A, args.B, D, args.maybe_group_scales, args.maybe_group_zeros, - args.maybe_group_size, args.maybe_channel_scales, - args.maybe_token_scales); - TORCH_CHECK(MacheteKernel::can_implement(arguments), - "Machete kernel cannot be run with these arguments"); - - size_t workspace_size = MacheteKernel::get_workspace_size(arguments); - torch::Tensor workspace = torch::empty( - workspace_size, torch::TensorOptions().dtype(torch::kU8).device(device)); - - MacheteKernel::run(arguments, workspace.mutable_data_ptr(), stream); - - return D; -}; - -}; // namespace machete \ No newline at end of file diff --git a/csrc/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu deleted file mode 100644 index 05a51ee21dd..00000000000 --- a/csrc/quantization/machete/machete_pytorch.cu +++ /dev/null @@ -1,73 +0,0 @@ -#include "machete_mm_launcher.cuh" -#include "machete_prepack_launcher.cuh" -#include "core/scalar_type.hpp" - -#include "core/registration.h" - -namespace machete { - -using namespace vllm; - -std::vector supported_schedules( - at::ScalarType a_type, int64_t b_type_id, - std::optional maybe_group_scales_type, - std::optional maybe_group_zeros_type, - std::optional maybe_channel_scales_type, - std::optional maybe_token_scales_type, - std::optional maybe_out_type) { - ScalarType const b_type = ScalarType::from_id(b_type_id); - return supported_schedules_dispatch({ - .a_type = a_type, - .b_type = b_type, - .maybe_group_scales_type = maybe_group_scales_type, - .maybe_group_zeros_type = maybe_group_zeros_type, - .maybe_channel_scales_type = maybe_channel_scales_type, - .maybe_token_scales_type = maybe_token_scales_type, - .maybe_out_type = maybe_out_type, - }); -} - -torch::Tensor mm(torch::Tensor const& A, torch::Tensor const& B, - int64_t b_type_id, - std::optional const& maybe_out_type, - std::optional const& maybe_group_scales, - std::optional const& maybe_group_zeros, - std::optional maybe_group_size, - std::optional const& maybe_channel_scales, - std::optional const& maybe_token_scales, - std::optional maybe_schedule) { - ScalarType const b_type = ScalarType::from_id(b_type_id); - return mm_dispatch({.A = A, - .B = B, - .b_type = b_type, - .maybe_out_type = maybe_out_type, - .maybe_group_scales = maybe_group_scales, - .maybe_group_zeros = maybe_group_zeros, - .maybe_group_size = maybe_group_size, - .maybe_channel_scales = maybe_channel_scales, - .maybe_token_scales = maybe_token_scales, - .maybe_schedule = maybe_schedule}); -} - -torch::Tensor prepack_B( - torch::Tensor const& B, at::ScalarType const& a_type, int64_t b_type_id, - std::optional const& maybe_group_scales_type) { - ScalarType const b_type = ScalarType::from_id(b_type_id); - return prepack_B_dispatch( - {.B = B, - .a_type = a_type, - .b_type = b_type, - .maybe_group_scales_type = maybe_group_scales_type}); -} - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("machete_prepack_B", &prepack_B); - m.impl("machete_mm", &mm); -} - -// use CatchAll since supported_schedules has no tensor arguments -TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) { - m.impl("machete_supported_schedules", &supported_schedules); -} - -}; // namespace machete diff --git a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu deleted file mode 100644 index 7d4c97fb57e..00000000000 --- a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ /dev/null @@ -1,106 +0,0 @@ - - -#include "marlin.cuh" - -#include "core/registration.h" - -// for only non-zp format (like gptq) -__global__ void marlin_int4_fp8_preprocess_kernel_without_zp( - // qweight: (size_k * size_n // 8,) - const int32_t* __restrict__ qweight, - // output: same shape with qweight - int32_t* __restrict__ output) { - int32_t val = qweight[blockIdx.x * 32 + threadIdx.x]; - int32_t new_val = 0; - -#pragma unroll - for (int32_t i = 0; i < 8; i++) { - int32_t single_val = val & 0xF; - single_val = single_val >= 8 ? single_val - 8 : 15 - single_val; - new_val |= single_val << (i * 4); - val >>= 4; - } - - output[blockIdx.x * 32 + threadIdx.x] = new_val; -} - -// for awq format only (with zp and with awq weight layout) -__global__ void marlin_int4_fp8_preprocess_kernel_awq( - // AWQ qweight: (size_k, size_n // 8) - const int32_t* __restrict__ qweight, - // output: same shape with qweight - int32_t* __restrict__ output, - // AWQ zeros: (size_k // group_size, size_n // 8) - const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, - int32_t group_size) { - int32_t val = - qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; - int32_t zero = - qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + - blockIdx.y]; - int32_t new_val = 0; - -#pragma unroll - for (int32_t i = 0; i < 8; i++) { - int32_t single_val = val & 0xF; - int32_t single_zero = zero & 0xF; - - single_val = - single_val >= single_zero ? single_val - single_zero : 15 - single_val; - new_val |= single_val << (i * 4); - val >>= 4; - zero >>= 4; - } - - output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; -} - -torch::Tensor marlin_int4_fp8_preprocess( - torch::Tensor& qweight, std::optional qzeros_or_none, - bool inplace) { - TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU"); - TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int, - "qweight.dtype != torch.int32"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - - torch::Tensor output = inplace ? qweight : torch::empty_like(qweight); - - if (!qzeros_or_none.has_value()) { - TORCH_CHECK(qweight.numel() * 8 % 256 == 0, - "qweight.numel() * 8 % 256 != 0"); - - int blocks = qweight.numel() * 8 / 256; - marlin_int4_fp8_preprocess_kernel_without_zp<<>>( - (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr()); - } else { - int32_t size_k = qweight.size(0); - int32_t size_n = qweight.size(1) * 8; - torch::Tensor qzeros = qzeros_or_none.value(); - - TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0"); - TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU"); - TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int, - "qweight.dtype != torch.int32"); - TORCH_CHECK(device_of(qweight) == device_of(qzeros), - "qzeros is not on the same device with qweight"); - - int32_t group_size = qweight.size(0) / qzeros.size(0); - TORCH_CHECK(qweight.size(1) == qzeros.size(1), - "qweight.size(1) != qzeros.size(1)"); - TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0, - "qweight.size(0) % qzeros.size(0) != 0"); - TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); - - dim3 blocks(size_k / 32, size_n / 8); - marlin_int4_fp8_preprocess_kernel_awq<<>>( - (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(), - (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size); - } - - return output; -} - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess); -} diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index a2170e48320..6c3456d06f2 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -283,6 +283,29 @@ __quickreduce_device_inline__ int packed_rcp(int a) { return R.i; } +template +__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low, + int16_t high); + +template <> +__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low, + int16_t high) { + // Convert two signed integers to one fp16x2 packed 32-bit lane. + half2 h = __halves2half2(__int2half_rn(static_cast(low)), + __int2half_rn(static_cast(high))); + return __builtin_bit_cast(int, h); +} + +template <> +__quickreduce_device_inline__ int packed_from_int16_pair( + int16_t low, int16_t high) { + // Convert two signed integers to one bf16x2 packed 32-bit lane. + nv_bfloat16 bf_low = __float2bfloat16(static_cast(low)); + nv_bfloat16 bf_high = __float2bfloat16(static_cast(high)); + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high); + return *reinterpret_cast(&bf2); +} + // changes dtype __quickreduce_device_inline__ float T2float_cast(half a) { return __half2float(a); diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/quickreduce/quick_reduce.h index 4cc35300bf8..7506329972b 100644 --- a/csrc/quickreduce/quick_reduce.h +++ b/csrc/quickreduce/quick_reduce.h @@ -59,11 +59,30 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, flag_color, this->kMaxProblemSize); \ } +// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8 +// the 3-bit codec's pack/unpack overhead outweighs the reduced communication +// volume, so INT3 is restricted to a TP2-only dispatch here. +#define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \ + if (world_size == 2) { \ + using LineCodec = __codec; \ + using AllReduceKernel = AllReduceTwoshot; \ + hipLaunchKernelGGL((allreduce_prototype_twoshot), \ + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ + num_blocks, rank, dbuffer_list, data_offset, \ + flag_color, this->kMaxProblemSize); \ + } else { \ + throw std::runtime_error( \ + "INT3 quick all-reduce is only supported for world_size == 2 " \ + "(TP2); use INT4/NONE for larger world sizes."); \ + } + enum QuickReduceQuantLevel { - F16 = 0, - INT8 = 1, - INT6 = 2, - INT4 = 3, + // Keep these ids in sync with Python QuickReduceRegime enum. + F16 = 0, // full-precision fp16/bf16 communication + INT8 = 1, // symmetric int8 + per-block scale + INT6 = 2, // symmetric int6 + per-block scale + INT4 = 3, // symmetric int4 + per-block scale + INT3 = 4, // symmetric int3 + per-block scale (TP2 only) }; struct DeviceComms { @@ -184,6 +203,9 @@ struct DeviceComms { case QuickReduceQuantLevel::INT4: TWOSHOT_DISPATCH(CodecQ4) break; + case QuickReduceQuantLevel::INT3: + TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3) + break; default: TWOSHOT_DISPATCH(CodecFP) break; diff --git a/csrc/quickreduce/quick_reduce_impl.cuh b/csrc/quickreduce/quick_reduce_impl.cuh index 38dc9938fc8..e9586e40ff0 100644 --- a/csrc/quickreduce/quick_reduce_impl.cuh +++ b/csrc/quickreduce/quick_reduce_impl.cuh @@ -2,6 +2,7 @@ #include #include "base.h" +#include namespace quickreduce { @@ -206,6 +207,168 @@ struct CodecQ4 : public CodecBase { } }; +// Int3 symmetric quantization codec. +// We quantize the FP16 data to block-scaled Int3 in blocks of 4 * +// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip), +// matching the structure of CodecQ4. Signed range is [-4, +3]. +template +struct CodecQ3 : public CodecBase { + static constexpr int kWorldSize = world_size; + + // Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes): + // - each thread owns 8 values and writes: + // * q2 payload : 8 * 2 bits -> uint16 (2 bytes) + // * q1 payload : 8 * 1 bit -> uint8 (1 byte) + // - one scale is shared per 32 values and written by group leader. + // + // kRankTileStride is split as: + // [0 .. 511] : q2 payload region (256 threads * 2 bytes) + // [512 .. 767] : q1 payload region (256 threads * 1 byte) + // [768 .. 895] : scale region (32 groups * 4 bytes) + static constexpr int kRankAtoms = kAtoms / kWorldSize; + static constexpr int kRankTileStride = 896; + static constexpr int kRankTileQ1Offset = 512; + static constexpr int kRankTileScaleOffset = 768; + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms; + static_assert(kRankTransmittedTileSize % 16 == 0, + "kRankTransmittedTileSize must be 16B aligned."); + + static constexpr int kRankBufferTileStride = + kRankTileStride / sizeof(int32x4_t); + + static constexpr int kTransmittedTileSize = + kRankTransmittedTileSize * kWorldSize; + + // {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps + // to -4; the sign cancels with decoding_scale on the recv side. + static constexpr int kScaleFactor = + std::is_same::value ? 0xB400B400 : 0xBE80BE80; + + // {1e-7, 1e-7}, f16x2_t + static constexpr int kScaleEpsilon = + std::is_same::value ? 0x00010001 : 0x33D733D7; + + // {-4, -4}, f16x2_t / bf16x2_t + static constexpr int kRangeMin = + std::is_same::value ? 0xC400C400 : 0xC080C080; + + // {+3, +3}, f16x2_t / bf16x2_t + static constexpr int kRangeMax = + std::is_same::value ? 0x42004200 : 0x40404040; + + // {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7]. + static constexpr int kRangeBias = 0x00040004; + + __quickreduce_device_inline__ CodecQ3(int thread, int rank) + : CodecBase(thread, rank) {} + + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer, + const int32x4_t* __restrict__ data) { + for (int k = 0; k < kRankAtoms; k++) { + int32x4_t const atom = data[k]; + + // 1) Per-group dynamic scale (shared across 32 values). + int wblockmax = group_abs_max(atom); + int decoding_scale = packed_mul(wblockmax, kScaleFactor); + int encoding_scale = packed_add(decoding_scale, kScaleEpsilon); + encoding_scale = packed_rcp(encoding_scale); + + // 2) Scale + clip to signed int3 range [-4, +3]. + int32x4_t w; + for (int i = 0; i < 4; i++) { + w[i] = packed_mul(atom[i], encoding_scale); + w[i] = packed_max(w[i], kRangeMin); + w[i] = packed_min(w[i], kRangeMax); + } + + // 3) Round to integer and bias to unsigned domain [0, 7]. + int32x4_t q; + { + int16_t* qi = reinterpret_cast(&q); + T* wh = reinterpret_cast(&w); + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i])); + + for (int i = 0; i < 4; i++) { + q[i] = packed_add(q[i], kRangeBias); + } + } + + // 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit + // halves, packed into one uint16 (low 2 bits per value) plus one + // uint8 (high 1 bit per value). + uint16_t q2w = 0; + uint8_t q1w = 0; + { + int16_t* tw = reinterpret_cast(&q); +#pragma unroll + for (int i = 0; i < 8; i++) { + uint32_t v = static_cast(tw[i]) & 0x7u; + q2w |= static_cast((v & 0x3u) << (i * 2)); + q1w |= static_cast(((v >> 2) & 0x1u) << i); + } + } + + uint8_t* atom_ptr = + reinterpret_cast(send_buffer + k * kRankBufferTileStride); + uint16_t* q2w_ptr = reinterpret_cast(atom_ptr) + thread; + uint8_t* q1w_ptr = + reinterpret_cast(atom_ptr + kRankTileQ1Offset) + thread; + int* qs_ptr = reinterpret_cast(atom_ptr + kRankTileScaleOffset) + + (thread / 8); + + __builtin_nontemporal_store(q2w, q2w_ptr); + *q1w_ptr = q1w; + if (threadIdx.x == group_leader) { + __builtin_nontemporal_store(decoding_scale, qs_ptr); + } + } + } + + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer, + int32x4_t* __restrict__ data) { + for (int k = 0; k < kRankAtoms; k++) { + uint8_t* atom_ptr = reinterpret_cast(*recv_buffer); + uint16_t* q2w_ptr = reinterpret_cast(atom_ptr) + thread; + uint8_t* q1w_ptr = + reinterpret_cast(atom_ptr + kRankTileQ1Offset) + thread; + int* qs_ptr = reinterpret_cast(atom_ptr + kRankTileScaleOffset) + + (thread / 8); + + uint16_t q2w = __builtin_nontemporal_load(q2w_ptr); + uint8_t q1w = *q1w_ptr; + int qs = __builtin_nontemporal_load(qs_ptr); + + *recv_buffer += kRankBufferTileStride; + + // Unpack unsigned values [0, 7] then shift back to signed domain + // [-4, +3] by adding kRangeMin. + int32x4_t w; + { + int16_t qv[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { + uint32_t low2 = (q2w >> (2 * i)) & 0x3u; + uint32_t high1 = (q1w >> i) & 0x1u; + qv[i] = static_cast(low2 | (high1 << 2)); + } + +#pragma unroll + for (int i = 0; i < 4; i++) { + int qpack = packed_from_int16_pair(qv[2 * i], qv[2 * i + 1]); + w[i] = packed_add(qpack, kRangeMin); + } + } + + // Apply decode scale to reconstruct fp16/bf16 lanes. + for (int i = 0; i < 4; i++) { + w[i] = packed_mul(w[i], qs); + } + + data[k] = w; + } + } +}; + // Int6 symmetric quantization codec. // We quantize the FP16 data to block-scaled Int6 in blocks of 4 * // kThreadGroupSize. @@ -377,7 +540,6 @@ struct CodecQ6 : public CodecBase { w[i] = packed_mul(w[i], qs); } - // That's pretty much it... data[k] = w; } } diff --git a/csrc/qutlass_registration.cpp b/csrc/qutlass_registration.cpp new file mode 100644 index 00000000000..effb4404135 --- /dev/null +++ b/csrc/qutlass_registration.cpp @@ -0,0 +1,5 @@ +#include "core/registration.h" + +// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp. +// This stub lets Python import vllm._qutlass_C to trigger op registration. +REGISTER_EXTENSION(_qutlass_C) diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 9e6c0726d19..4ac255d0a75 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* q_ptr = q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE; const _B16x8* q_ptrh8 = reinterpret_cast(q_ptr); - const int qhead_elemh8 = laneid / 4; + const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1); for (int h = 0; h < QHLOOP - 1; h++) { const int qhead_idx = h * 4 + lane4id; diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 10e3cbf2e0b..615cdabed58 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -70,6 +70,15 @@ bool on_gfx12() { return result; } +bool on_gfx1151() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx1151") != std::string::npos; + }(); + return result; +} + #if defined(NDEBUG) #undef NDEBUG #include @@ -1237,6 +1246,45 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ } +// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N). +// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9) +// WVPRGRP= waves per group (always 16) +// YTILE = output rows per thread tile +// UNRL = K-loop unroll factor +// N = batch size (passed through from the switch in wvSplitK) +#define WVSPLIT_TILE(_sYT, __N) \ + { \ + if (on_gfx1151()) { \ + bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ + if (_sYT <= 1) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if ((K_in % 1024 == 512) && K_in >= 1536 && \ + (_sYT >= 40 || K_in >= 4096)) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/4, /*UNRL=*/1, \ + __N) \ + else if (K_in < 1024) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/2, /*UNRL=*/4, \ + __N) \ + else if (K_in <= 2048 && (__N >= 2 || _sYT <= 26)) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if (__N >= 2 && !fit_lds) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if (__N == 1) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/2, \ + __N) \ + else \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/1, \ + __N) \ + } else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */ \ + WVSPLIT_TILE_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, _sYT, __N) \ + } else { /* GFX9, wave64 */ \ + WVSPLIT_TILE_CFG(/*THRDS=*/64, /*WVPRGRP=*/16, _sYT, __N) \ + } \ + } + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] { using fptype = typename scalar::type; fptype* af4 = reinterpret_cast(in_a.data_ptr()); @@ -1251,37 +1299,21 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, // then cut the active waves to balance their distribution... int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); - const bool use_wave32 = on_gfx1x(); switch (N_in) { case 1: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 1) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 1) + WVSPLIT_TILE(sYT, 1) break; case 2: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 2) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 2) + WVSPLIT_TILE(sYT, 2) break; case 3: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 3) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 3) + WVSPLIT_TILE(sYT, 3) break; case 4: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 4) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 4) + WVSPLIT_TILE(sYT, 4) break; case 5: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 5) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 5) + WVSPLIT_TILE(sYT, 5) break; default: throw std::runtime_error( diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index c63e59c3b03..bcf0ce3e1c4 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -2,7 +2,6 @@ // cache.h, which is no longer included here after cache ops moved to // _C_stable_libtorch). #include -#include "cuda_utils.h" #include "ops.h" #include "core/registration.h" #include @@ -18,140 +17,6 @@ // https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations -TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { - // vLLM custom ops - // - - ops.def( - "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " - "y_q, Tensor! y_s," - "bool use_ue8m0) -> ()"); - ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA, - &persistent_masked_m_silu_mul_quant); - - ops.def("weak_ref_tensor(Tensor input) -> Tensor"); - ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor); - -#ifdef USE_ROCM - // TODO: Remove this once we upgrade to torch 2.11. - // ROCm still uses torch 2.10, - // So we still need to use unstable torch ABI for now. - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU, - &get_cuda_view_from_cpu_tensor); -#endif - - // Activation ops (quantized only — basic ops moved to _C_stable_libtorch) - ops.def( - "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); - ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant); - - // 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. Registered in _C_stable_libtorch (incl. the FlashInfer V4 - // full-cache bf16/fp8 variants). - - // Quantization ops -#ifndef USE_ROCM - - // Note about marlin kernel 'workspace' arguments: - // Technically these should be mutable since they are modified by the kernel. - // But since they are set back to zero once the kernel is finished we can - // hand wave and say that they have no net effect. - // - // The reason to mark 'workspace' as immutable is so that they don't interfere - // with using ScalarType arguments in the ops. If they are marked as mutable, - // pytorch throws an assert in - // 'torch._higher_order_ops._register_effectful_op' that prevents these - // kernels from being torch.compile'd. - // See the following document for more info on custom types and ops that use - // custom types: - // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - - // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. - ops.def( - "machete_supported_schedules(" - " ScalarType a_type," - " int b_type," - " ScalarType? maybe_group_scales_type," - " ScalarType? maybe_group_zeros_type," - " ScalarType? maybe_channel_scales_type," - " ScalarType? maybe_token_scales_type," - " ScalarType? maybe_out_type" - ") -> str[]"); - ops.def( - "machete_mm(" - " Tensor A," - " Tensor B," - " int b_type," - " ScalarType? out_type," - " Tensor? group_scales," - " Tensor? group_zeros," - " int? group_size," - " Tensor? channel_scales," - " Tensor? token_scales," - " str? schedule" - ") -> Tensor"); - ops.def( - "machete_prepack_B(" - " Tensor B," - " ScalarType a_type," - " int b_type," - " ScalarType? group_scales_type" - ") -> Tensor"); - // conditionally compiled so impl registration is in source file - - // Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4). - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor? " - "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registration is in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file - -#endif - -#ifndef USE_ROCM - // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+). - ops.def( - "mxfp8_experts_quant(" - " Tensor input, Tensor problem_sizes, Tensor expert_offsets," - " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)" - " -> ()"); - // conditionally compiled so impl registration is in source file - - // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+). - ops.def( - "cutlass_mxfp8_grouped_mm(" - " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out," - " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)" - " -> ()"); - // conditionally compiled so impl registration is in source file - -#endif -} - #ifdef USE_ROCM TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). @@ -169,20 +34,6 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { custom_ar.def("qr_max_size", &qr_max_size); } - -// TODO: Remove this once ROCm upgrade to torch 2.11. -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { - // Cuda utils - // Gets the specified device attribute. - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.impl("get_device_attribute", &get_device_attribute); - - // Gets the maximum shared memory per block device attribute. - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - &get_max_shared_memory_per_block_device_attribute); -} #endif REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/docker/Dockerfile b/docker/Dockerfile index 9b4227cdf65..b47853a06c7 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -99,9 +99,16 @@ ARG INSTALL_KV_CONNECTORS=false # prepare basic build environment FROM ${BUILD_BASE_IMAGE} AS base +ARG TARGETPLATFORM ARG CUDA_VERSION ARG PYTHON_VERSION ARG BUILD_OS +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 ENV DEBIAN_FRONTEND=noninteractive @@ -148,11 +155,13 @@ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ sudo \ python3-pip \ libibverbs-dev \ - # Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519 - # as it was causing spam when compiling the CUTLASS kernels - gcc-10 \ - g++-10 \ - && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \ + # GCC 10 was previously pinned to suppress spurious -Wredundant-move warnings + # from CUTLASS (https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519). That bug + # was fixed in GCC 11. GCC >= 11.3 is now required because PyTorch's C++20 headers + # (pytorch/pytorch#167929) are not compatible with GCC < 11.3. + gcc-11 \ + g++-11 \ + && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 --slave /usr/bin/g++ g++ /usr/bin/g++-11 \ # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04 # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA) @@ -160,6 +169,27 @@ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ && rm -rf /var/lib/apt/lists/*; \ fi +# Install sccache once in base so Rust and CMake/CUDA build stages share the +# same binary and remote cache configuration. +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Installing sccache..." \ + && case "${TARGETPLATFORM}" in \ + linux/arm64) SCCACHE_ARCH="aarch64" ;; \ + linux/amd64) SCCACHE_ARCH="x86_64" ;; \ + *) echo "Unsupported TARGETPLATFORM for sccache: ${TARGETPLATFORM}" >&2; exit 1 ;; \ + esac \ + && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \ + && tar -xzf sccache.tar.gz \ + && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl; \ + fi + +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + # Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all # downstream stages stay distro-agnostic. RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" "${UV_INSTALL_DIR}" \ @@ -218,6 +248,10 @@ COPY requirements/common.txt requirements/common.txt COPY requirements/cuda.txt requirements/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY pyproject.toml pyproject.toml +# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that +# share paths with different content. uv can extract them in either order, +# leaving base files that break CUDA 13 CuTe DSL JIT. +# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. RUN --mount=type=cache,target=/opt/uv/cache \ if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \ sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \ @@ -234,6 +268,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + fi \ + && if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ + CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ + if [ -n "$CUTLASS_DSL_VERSION" ]; then \ + uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \ + "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ + fi; \ fi # Track PyTorch lib versions used during build and match in downstream instances. @@ -248,63 +289,78 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Explicitly set the list to avoid issues with torch 2.2 # See https://github.com/pytorch/pytorch/pull/123243 # From versions.json: .torch.cuda_arch_list -ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' +# Do not add +PTX here: vLLM filters torch's top-level PTX flag when it +# converts global gencode flags into per-kernel arch lists. If a specific +# kernel needs PTX, add +PTX to that kernel's CMake arch list instead. +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} #################### BUILD BASE IMAGE #################### #################### RUST BUILD IMAGE #################### # Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel # build stage doesn't need the rust toolchain, protoc, or the rust source. -# This stage runs in parallel with csrc-build/extensions-build. -FROM ${BUILD_BASE_IMAGE} AS rust-build +# This stage reuses the Python environment from base and runs in parallel with +# csrc-build/extensions-build. +FROM base AS rust-build ARG BUILD_OS +ARG USE_SCCACHE +ARG SCCACHE_ENDPOINT -ENV DEBIAN_FRONTEND=noninteractive - -# Install a basic C toolchain (some rust crates compile C in their build.rs -# scripts) and unzip (used to extract the pinned protoc release below). +# Install native tools needed only for Rust/protoc builds. RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ dnf install -y --setopt=install_weak_deps=False \ - ca-certificates curl git gcc gcc-c++ make unzip \ + make unzip \ && dnf clean all && rm -rf /var/cache/dnf; \ else \ apt-get update -y \ && apt-get install -y --no-install-recommends \ - ca-certificates curl git build-essential unzip \ + make unzip \ && rm -rf /var/lib/apt/lists/*; \ fi COPY tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" - WORKDIR /workspace -# Copy only the rust workspace — the binary is the sole artifact we need. +COPY requirements/build/rust.txt requirements/build/rust.txt +RUN --mount=type=cache,target=/opt/uv/cache \ + uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt + +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml +COPY tools/build_rust.py tools/build_rust.py COPY build_rust.sh build_rust.sh # Cap cargo parallelism to avoid exhausting the CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -# Build the release binary. Cache cargo registry/git and target/, but copy the -# binary out of the target/ cache mount so it persists into the image layer -# for later COPY --from=rust-build. -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ - --mount=type=cache,target=/workspace/rust/target \ - VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh +# BuildKit can run this stage in parallel with csrc-build. Keep Rust on a +# separate local sccache daemon while sharing the same remote cache backend. +ENV SCCACHE_SERVER_PORT=4227 + +# Build the release artifacts. Cache cargo registry/git, but not target/, +# because stale target metadata can outlive source updates across BuildKit +# cache reuse. +RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ + if [ "$USE_SCCACHE" = "1" ]; then \ + if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \ + export RUSTC_WRAPPER=sccache; \ + sccache --show-stats; \ + fi \ + && bash build_rust.sh \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + sccache --show-stats; \ + fi #################### RUST BUILD IMAGE #################### #################### CSRC BUILD IMAGE #################### FROM base AS csrc-build -ARG TARGETPLATFORM ARG PIP_INDEX_URL UV_INDEX_URL ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL @@ -342,6 +398,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ WORKDIR /workspace COPY pyproject.toml setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py COPY cmake cmake/ COPY csrc csrc/ COPY vllm/envs.py vllm/envs.py @@ -355,11 +412,7 @@ ARG nvcc_threads=8 ENV NVCC_THREADS=$nvcc_threads ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 # Flag to control whether to use pre-built vLLM wheels ARG VLLM_USE_PRECOMPILED="" @@ -389,22 +442,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ RUN --mount=type=cache,target=/opt/uv/cache \ --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && case "${TARGETPLATFORM}" in \ - linux/arm64) SCCACHE_ARCH="aarch64" ;; \ - linux/amd64) SCCACHE_ARCH="x86_64" ;; \ - *) echo "Unsupported TARGETPLATFORM for sccache: ${TARGETPLATFORM}" >&2; exit 1 ;; \ - esac \ - && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \ - && tar -xzf sccache.tar.gz \ - && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl \ - && if [ ! -z ${SCCACHE_ENDPOINT} ] ; then export SCCACHE_ENDPOINT=${SCCACHE_ENDPOINT} ; fi \ - && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \ - && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \ - && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \ - && export SCCACHE_IDLE_TIMEOUT=0 \ + if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi \ && export CMAKE_BUILD_TYPE=Release \ && export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" \ && export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \ @@ -506,9 +544,10 @@ WORKDIR /workspace COPY --from=csrc-build /workspace/dist /precompiled-wheels COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ @@ -535,9 +574,17 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 +# Record the wheel checksum so downstream stages can bust their layer cache +# when the wheel changes, without copying the wheel itself into the image. +RUN sha256sum dist/*.whl > dist/wheel.sha256 + # Copy extension wheels from extensions-build stage for later use COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist +# Record the EP kernels wheel checksum for the same cache-busting purpose. +RUN sha256sum /tmp/ep_kernels_workspace/dist/*.whl \ + > /tmp/ep_kernels_workspace/dist/wheels.sha256 + # Check the size of the wheel if RUN_WHEEL_CHECK is true COPY .buildkite/check-wheel-size.py check-wheel-size.py # sync the default value with .buildkite/check-wheel-size.py @@ -745,6 +792,10 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0 ARG PYTORCH_CUDA_INDEX_BASE_URL COPY requirements/common.txt /tmp/common.txt COPY requirements/cuda.txt /tmp/requirements-cuda.txt +# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that +# share paths with different content. uv can extract them in either order, +# leaving base files that break CUDA 13 CuTe DSL JIT. +# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. RUN --mount=type=cache,target=/opt/uv/cache \ if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \ sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \ @@ -752,12 +803,19 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ uv pip install --system -r /tmp/requirements-cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \ + if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ + CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ + if [ -n "$CUTLASS_DSL_VERSION" ]; then \ + uv pip install --system --force-reinstall --no-deps \ + "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ + fi; \ + fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.12 +ARG FLASHINFER_VERSION=0.6.13 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') @@ -796,7 +854,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \ fi; \ - uv pip install --system accelerate modelscope \ + uv pip install --system accelerate 'modelscope<1.38' \ "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}" # ============================================================ @@ -814,6 +872,11 @@ ARG PYTORCH_NIGHTLY # Install vLLM wheel first, so that torch etc will be installed. # Check whether to install torch nightly instead of release for this build. COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt +# Copy only the wheel checksum (a few bytes) so a wheel change invalidates this +# install layer. The wheel itself is bind-mounted below and never enters the +# image. Without this the bind mount is not part of the layer cache key, so a +# warm BuildKit agent can skip the install and ship a stale wheel. +COPY --from=build /workspace/dist/wheel.sha256 /tmp/vllm-wheel.sha256 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \ --mount=type=cache,target=/opt/uv/cache \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ @@ -836,18 +899,27 @@ uv pip list # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH -# Install EP kernels wheels (DeepEP) that have been built in the `build` stage +# Install EP kernels wheels (DeepEP) that have been built in the `build` stage. +# As with the vLLM wheel above, copy only the checksum to bust the layer cache +# and bind-mount the wheel for the actual install to keep it out of the image. +COPY --from=build /tmp/ep_kernels_workspace/dist/wheels.sha256 /tmp/ep-kernels-wheels.sha256 RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \ --mount=type=cache,target=/opt/uv/cache \ uv pip install --system ep_kernels/dist/*.whl --verbose \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') -# Download FlashInfer precompiled cubins AFTER all pip installs are done. -# This must run after the vLLM wheel and EP kernels installs above, because -# those can reinstall/touch flashinfer packages. Downloading cubins earlier -# (in the flashinfer-jit-cache layer) causes ~2.5 GB of layer duplication -# when a later pip install overwrites flashinfer package files. -RUN flashinfer show-config && flashinfer download-cubin +# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that +# share paths with different content. Force -libs-cu13 last after runtime +# dependency installs so uv cannot leave base files behind. +# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. +RUN --mount=type=cache,target=/opt/uv/cache \ + if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ + CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ + if [ -n "$CUTLASS_DSL_VERSION" ]; then \ + uv pip install --system --force-reinstall --no-deps \ + "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ + fi; \ + fi # CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will # return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers @@ -957,7 +1029,8 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL ENV UV_HTTP_TIMEOUT=500 # install kv_connectors if requested -ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX' +# Do not add +PTX here; see the main TORCH_CUDA_ARCH_LIST comment above. +ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0' ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} RUN --mount=type=cache,target=/opt/uv/cache \ --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 02af3fd3c39..a528ffbd8d1 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common WORKDIR /workspace ARG PYTHON_VERSION=3.12 -ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" ARG max_jobs=32 ENV MAX_JOBS=${max_jobs} @@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH" ENV UV_HTTP_TIMEOUT=500 # Install Python dependencies -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" @@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --upgrade pip && \ - uv pip install -r requirements/cpu.txt + uv pip install -r requirements/cpu.txt --torch-backend cpu ARG TARGETARCH ENV TARGETARCH=${TARGETARCH} @@ -93,35 +90,34 @@ ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ - ca-certificates curl git build-essential unzip \ + ca-certificates curl git build-essential unzip python3 python3-pip \ && rm -rf /var/lib/apt/lists/* COPY tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" - WORKDIR /workspace -# Copy only the rust workspace — the binary is the sole artifact we need. +COPY requirements/build/rust.txt requirements/build/rust.txt +RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt + +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml +COPY tools/build_rust.py tools/build_rust.py COPY build_rust.sh build_rust.sh # Cap cargo parallelism to avoid exhausting the CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -# Build the release binary. Cache cargo registry/git and target/, but copy the -# binary out of the target/ cache mount so it persists into the image layer -# for later COPY --from=rust-build. +# Build the release artifacts. Cache cargo registry/git, but not target/, +# because stale target metadata can outlive source updates across BuildKit +# cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ - --mount=type=cache,target=/workspace/rust/target,sharing=locked \ - VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh + bash build_rust.sh ######################### BUILD IMAGE ######################### FROM base AS vllm-build @@ -150,13 +146,14 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \ COPY requirements/build/cpu.txt requirements/build/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/build/cpu.txt + uv pip install -r requirements/build/cpu.txt --torch-backend cpu COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi @@ -168,6 +165,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ######################### TRITON-CPU BUILD IMAGE ######################### FROM base AS vllm-triton-cpu-build +# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ... +# Re-declared here because this stage is `FROM base` (not `vllm-build`), so it +# does not inherit the ARG/ENV defined there. Without it, the guard below would +# see an empty value and build triton-cpu on non-x86 targets (e.g. arm64). +ARG VLLM_CPU_X86=0 + WORKDIR /vllm-workspace RUN mkdir dist @@ -187,29 +190,19 @@ FROM base AS vllm-test-deps WORKDIR /vllm-workspace -# Copy test requirements -COPY requirements/test/cuda.in requirements/test/cpu.in +# Test requirements are compiled from requirements/test/cuda.in into +# requirements/test/cpu.txt by the pip-compile-cpu pre-commit hook, which +# resolves CPU wheels via uv's --torch-backend cpu. +COPY requirements/test/cpu.txt requirements/test/cpu.txt -RUN \ - 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/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/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 +# cpu.txt is compiled for x86_64, so platform markers are resolved away. Drop +# packages unavailable on aarch64 (decord, terratorch) for arm builds. +RUN case "$(uname -m)" in \ + aarch64|arm64) sed -i '/^decord==/d; /^terratorch==/d' requirements/test/cpu.txt ;; \ + esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/test/cpu.txt + uv pip install -r requirements/test/cpu.txt --torch-backend cpu ######################### DEV IMAGE ######################### FROM vllm-build AS vllm-dev @@ -235,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements/lint.txt && \ - uv pip install -r requirements/test/cpu.txt && \ + uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \ pre-commit install --hook-type pre-commit --hook-type commit-msg ENTRYPOINT ["bash"] @@ -269,6 +262,11 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60 ######################### RELEASE IMAGE ######################### FROM base AS vllm-openai +# Re-declared here because this stage is `FROM base` (not `vllm-build`), so the +# RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86 would +# otherwise see an empty value and try to install it on non-x86 targets. +ARG VLLM_CPU_X86=0 + WORKDIR /vllm-workspace RUN --mount=type=cache,target=/root/.cache/uv \ @@ -300,6 +298,12 @@ LABEL ai.vllm.build.cpu-x86="${VLLM_CPU_X86:-false}" LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}" LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}" +# Copy the examples directory (including the chat/tool templates) so it is +# present in the released image, as the CUDA image ships it too. The vllm-test +# stage above adds examples/ for testing only, so without this the published +# vllm-openai-cpu image would not ship examples/*.jinja. +COPY examples examples + ENTRYPOINT ["vllm", "serve"] diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch deleted file mode 100644 index 4fbfe832ac3..00000000000 --- a/docker/Dockerfile.nightly_torch +++ /dev/null @@ -1,325 +0,0 @@ -####### -# -# THIS FILE IS DEPRECATED AND WILL BE REMOVED SHORTLY -# -# Please use the standard Dockerfile with PYTORCH_NIGHTLY=1 instead -# -####### - -# The vLLM Dockerfile is used to construct vLLM image against torch nightly that can be directly used for testing - -# for torch nightly, cuda >=12.6 is required, -# use 12.8 due to FlashAttention issue with cuda 12.6 (https://github.com/vllm-project/vllm/issues/15435#issuecomment-2775924628) -ARG CUDA_VERSION=12.8.0 -# -#################### BASE BUILD IMAGE #################### -# prepare basic build environment -FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS base -ARG CUDA_VERSION=12.8.0 -ARG PYTHON_VERSION=3.12 -ARG TARGETPLATFORM -ENV DEBIAN_FRONTEND=noninteractive -# Install Python and other dependencies -RUN apt-get update -y \ - && apt-get install -y ccache software-properties-common git curl sudo \ - && 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 \ - && apt-get update -y \ - && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - && 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 \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version \ - && python3 -m pip --version -# Install uv for faster pip installs -RUN --mount=type=cache,target=/root/.cache/uv \ - python3 -m pip install uv - -# 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 -ENV UV_HTTP_TIMEOUT=500 - -# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519 -# as it was causing spam when compiling the CUTLASS kernels -RUN apt-get install -y gcc-10 g++-10 -RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 -RUN < torch_build_versions.txt -RUN cat torch_build_versions.txt - -# cuda arch list used by torch -# can be useful for `test` -# explicitly set the list to avoid issues with torch 2.2 -# see https://github.com/pytorch/pytorch/pull/123243 - -#################### BASE BUILD IMAGE #################### - -#################### RUST BUILD IMAGE #################### -# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build -# stage doesn't need the rust toolchain or protoc. -FROM ubuntu:22.04 AS rust-build -ENV DEBIAN_FRONTEND=noninteractive - -RUN apt-get update -y \ - && apt-get install -y --no-install-recommends \ - ca-certificates curl git build-essential unzip \ - && rm -rf /var/lib/apt/lists/* - -COPY tools/install_protoc.sh /tmp/install_protoc.sh -RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh - -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" - -WORKDIR /workspace - -COPY rust rust -COPY rust-toolchain.toml rust-toolchain.toml -COPY build_rust.sh build_rust.sh - -# Cap cargo parallelism to avoid exhausting the CI host's open-file limit -# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). -ENV CARGO_BUILD_JOBS=4 - -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ - --mount=type=cache,target=/workspace/rust/target \ - VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh -#################### RUST BUILD IMAGE #################### - -#################### WHEEL BUILD IMAGE #################### -FROM base AS build -ARG TARGETPLATFORM - -# 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 -ENV UV_HTTP_TIMEOUT=500 - -COPY . . - -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs - -RUN python3 use_existing_torch.py - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/build/cuda.txt - -ARG GIT_REPO_CHECK=0 -RUN --mount=type=bind,source=.git,target=.git \ - if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi - -# Max jobs used by Ninja to build extensions -ARG max_jobs=16 -ENV MAX_JOBS=${max_jobs} -ARG nvcc_threads=2 -ENV NVCC_THREADS=$nvcc_threads - -ARG USE_SCCACHE -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - -# if USE_SCCACHE is set, use sccache to speed up compilation -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,source=.git,target=.git \ - if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \ - && tar -xzf sccache.tar.gz \ - && sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \ - && rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \ - && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \ - && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \ - && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \ - && export SCCACHE_IDLE_TIMEOUT=0 \ - && export CMAKE_BUILD_TYPE=Release \ - && sccache --show-stats \ - && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \ - && sccache --show-stats; \ - fi - -ENV CCACHE_DIR=/root/.cache/ccache -RUN --mount=type=cache,target=/root/.cache/ccache \ - --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,source=.git,target=.git \ - if [ "$USE_SCCACHE" != "1" ]; then \ - # Clean any existing CMake artifacts - rm -rf .deps && \ - mkdir -p .deps && \ - python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \ - fi - -#################### WHEEL BUILD IMAGE #################### - -################### VLLM INSTALLED IMAGE #################### -# Setup clean environment for vLLM and its dependencies for test and api server using ubuntu22.04 with AOT flashinfer -FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS vllm-base -# prepare for environment starts -ARG CUDA_VERSION=12.8.0 -ARG PYTHON_VERSION=3.12 -WORKDIR /vllm-workspace -ENV DEBIAN_FRONTEND=noninteractive -ARG TARGETPLATFORM - -RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \ - echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment - -# Install Python and other dependencies -RUN apt-get update -y \ - && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \ - && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \ - && 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 \ - && apt-get update -y \ - && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \ - && 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 \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version - -RUN --mount=type=cache,target=/root/.cache/uv \ - python3 -m pip install uv - -# 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 -ENV UV_HTTP_TIMEOUT=500 - -# Workaround for https://github.com/openai/triton/issues/2507 and -# https://github.com/pytorch/pytorch/issues/107960 -- hopefully -# this won't be needed for future versions of this docker image -# or future versions of triton. -RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/ - -# get the nightly torch version used in the build to make sure the version is the same -COPY --from=base /workspace/torch_build_versions.txt ./torch_build_versions.txt - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system $(cat torch_build_versions.txt | xargs) --index-url https://download.pytorch.org/whl/nightly/cu128 - -# install the vllm wheel -RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm-dist \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system vllm-dist/*.whl --verbose - -ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0' - -# install package for build flashinfer -# see issue: https://github.com/flashinfer-ai/flashinfer/issues/738 -RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.post1 - - -# build flashinfer for torch nightly from source around 10 mins -# release version: v0.6.12 -# 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.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \ - && cd flashinfer \ - && git submodule update --init --recursive \ - && echo "finish git clone flashinfer..." \ - && rm -rf build \ - && export TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} \ - && FLASHINFER_ENABLE_AOT=1 python3 setup.py bdist_wheel --dist-dir=../flashinfer-dist --verbose \ - && cd .. \ - && rm -rf flashinfer - -# install flashinfer -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system flashinfer-dist/*.whl --verbose - -# install common packages -COPY requirements/common.txt requirements/common.txt -COPY use_existing_torch.py use_existing_torch.py -COPY pyproject.toml pyproject.toml - -COPY examples examples -COPY benchmarks benchmarks -COPY ./vllm/collect_env.py . - -RUN python3 use_existing_torch.py -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/common.txt - -################### VLLM INSTALLED IMAGE #################### - - -#################### UNITTEST IMAGE ############################# -FROM vllm-base as test -COPY tests/ tests/ - -# install build and runtime dependencies without stable torch version -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 -ENV UV_HTTP_TIMEOUT=500 - -# install development dependencies (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -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 - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/test/nightly-torch.txt - -# Logging to confirm the torch versions -RUN pip freeze | grep -E 'torch|vllm|flashinfer' - -# Logging to confirm all the packages are installed -RUN pip freeze - -#################### UNITTEST IMAGE ############################# diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 845d900c39c..f0363d43be2 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -1,275 +1,80 @@ +# Base UBI image ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 ############################################################### -# Stage to build openblas +# BUILDER STAGE # ############################################################### -FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS openblas-builder - -ARG MAX_JOBS -ARG OPENBLAS_VERSION=0.3.30 -RUN microdnf install -y dnf && dnf install -y gcc-toolset-14 make wget unzip \ - && source /opt/rh/gcc-toolset-14/enable \ - && wget https://github.com/OpenMathLib/OpenBLAS/releases/download/v$OPENBLAS_VERSION/OpenBLAS-$OPENBLAS_VERSION.zip \ - && unzip OpenBLAS-$OPENBLAS_VERSION.zip \ - && cd OpenBLAS-$OPENBLAS_VERSION \ - && make -j${MAX_JOBS} TARGET=POWER9 BINARY=64 USE_OPENMP=1 USE_THREAD=1 NUM_THREADS=120 DYNAMIC_ARCH=1 INTERFACE64=0 \ - && cd /tmp && touch control - - -############################################################### -# base stage with dependencies coming from centos mirrors -############################################################### -FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS centos-deps-builder -RUN microdnf install -y dnf && \ - dnf install -y https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-gpg-keys-9.0-26.el9.noarch.rpm \ - https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-stream-repos-9.0-26.el9.noarch.rpm \ - https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ - dnf config-manager --set-enabled crb - -RUN dnf install -y openjpeg2-devel lcms2-devel tcl-devel tk-devel fribidi-devel yajl-devel && \ - dnf remove -y centos-gpg-keys-9.0-24.el9.noarch centos-stream-repos-9.0-26.el9.noarch - - -############################################################### -# base stage with basic dependencies -############################################################### - -FROM centos-deps-builder AS base-builder +FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS builder-base +ARG VLLM_VERSION="0.22.1" ARG PYTHON_VERSION=3.12 -ARG OPENBLAS_VERSION=0.3.30 - -# Set Environment Variables for venv, cargo & openblas -ENV VIRTUAL_ENV=/opt/vllm -ENV PATH=${VIRTUAL_ENV}/bin:/root/.cargo/bin:$PATH -ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig/ -ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib64:/usr/local/lib:/usr/lib64:/usr/lib -ENV UV_LINK_MODE=copy - -# install gcc-13, python, rust, openblas -# Note: A symlink for libatomic.so is created for gcc-13 (linker fails to find libatomic otherwise - reqd. for sentencepiece) -# Note: A dummy file 'control' is created in /tmp/ to artificially create dependencies between stages when building stages in parallel -# when `--jobs=` is passed with podman build command - -COPY --from=openblas-builder /tmp/control /dev/null - -RUN --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \ - dnf install -y openssl-devel \ - && dnf install -y \ - git tar gcc-toolset-14 automake libtool \ - pkgconfig xsimd zeromq-devel kmod findutils protobuf* \ - libtiff-devel libjpeg-devel zlib-devel freetype-devel libwebp-devel \ - harfbuzz-devel libraqm-devel libimagequant-devel libxcb-devel \ - python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip clang-devel \ - && dnf clean all \ - && PREFIX=/usr/local make -C /openblas install \ - && ln -sf /usr/lib64/libatomic.so.1 /usr/lib64/libatomic.so \ - && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ - && python -m pip install -U pip uv \ - && uv pip install wheel build "setuptools<70" setuptools_scm setuptools_rust meson-python 'cmake<4' ninja cython scikit_build_core scikit_build \ - && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ - && cd /tmp && touch control - - -############################################################### -# Stage to build torch family -############################################################### - -FROM base-builder AS torch-builder - -ARG MAX_JOBS -ARG TORCH_VERSION=2.7.0 -ARG _GLIBCXX_USE_CXX11_ABI=1 -ARG OPENBLAS_VERSION=0.3.30 - -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/pytorch.git -b v${TORCH_VERSION} && \ - cd pytorch && \ - uv pip install -r requirements.txt && \ - python setup.py develop && \ - rm -f dist/torch*+git*whl && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir /torchwheels/ - -ARG TORCHVISION_VERSION=0.22.0 -ARG TORCHVISION_USE_NVJPEG=0 -ARG TORCHVISION_USE_FFMPEG=0 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/vision.git -b v${TORCHVISION_VERSION} && \ - cd vision && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - BUILD_VERSION=${TORCHVISION_VERSION} \ - uv build --wheel --out-dir /torchwheels/ --no-build-isolation - -ARG TORCHAUDIO_VERSION=2.7.0 -ARG BUILD_SOX=1 -ARG BUILD_KALDI=1 -ARG BUILD_RNNT=1 -ARG USE_FFMPEG=0 -ARG USE_ROCM=0 -ARG USE_CUDA=0 -ARG TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} && \ - cd audio && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - BUILD_VERSION=${TORCHAUDIO_VERSION} \ - uv build --wheel --out-dir /torchwheels/ --no-build-isolation - -############################################################### -# Stage to build pyarrow -############################################################### - -FROM base-builder AS arrow-builder - -ARG MAX_JOBS -ARG PYARROW_PARALLEL -ARG PYARROW_VERSION=21.0.0 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/apache/arrow.git -b apache-arrow-${PYARROW_VERSION} && \ - cd arrow/cpp && \ - mkdir build && cd build && \ - cmake -DCMAKE_BUILD_TYPE=release \ - -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DARROW_PYTHON=ON \ - -DARROW_BUILD_TESTS=OFF \ - -DARROW_JEMALLOC=ON \ - -DARROW_BUILD_STATIC="OFF" \ - -DARROW_PARQUET=ON \ - .. && \ - make install -j ${MAX_JOBS:-$(nproc)} && \ - cd ../../python/ && \ - uv pip install -v -r requirements-build.txt && uv pip install numpy==2.1.3 && \ - PYARROW_PARALLEL=${PYARROW_PARALLEL:-$(nproc)} \ - python setup.py build_ext \ - --build-type=release --bundle-arrow-cpp \ - bdist_wheel --dist-dir /arrowwheels/ - -############################################################### -# Stage to build opencv -############################################################### - -FROM base-builder AS cv-builder - -ARG MAX_JOBS -ARG OPENCV_VERSION=86 -# patch for version 4.11.0.86 -ARG OPENCV_PATCH=97f3f39 -ARG ENABLE_HEADLESS=1 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/opencv/opencv-python.git -b ${OPENCV_VERSION} && \ - cd opencv-python && \ - sed -i -E -e 's/"setuptools.+",/"setuptools",/g' pyproject.toml && \ - cd opencv && git cherry-pick --no-commit $OPENCV_PATCH && cd .. && \ - uv pip install scikit-build && \ - python -m build --wheel --installer=uv --outdir /opencvwheels/ - -############################################################### -# Stage to build numactl -############################################################### - -FROM base-builder AS numa-builder - -# Note: Building numactl with gcc-11. Compiling with gcc-13 in this builder stage will -# trigger recompilation with gcc-11 (and require libtool) in the final stage where we do not have gcc-13 -ARG MAX_JOBS -ARG NUMACTL_VERSION=2.0.19 -RUN git clone --recursive https://github.com/numactl/numactl.git -b v${NUMACTL_VERSION} \ - && cd numactl \ - && autoreconf -i && ./configure \ - && make -j ${MAX_JOBS:-$(nproc)} - - -############################################################### -# Stage to build numba -############################################################### - -FROM base-builder AS numba-builder - -ARG MAX_JOBS -ARG NUMBA_VERSION=0.61.2 - -# Clone all required dependencies -RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-14/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \ - git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ - cd ./numba && \ - if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \ - sed -i '/#include "internal\/pycore_atomic.h"/i\#include "dynamic_annotations.h"' numba/_dispatcher.cpp; \ - fi && python -m build --wheel --installer=uv --outdir /numbawheels/ - -############################################################### -# Stage to build vllm - this stage builds and installs -# vllm, tensorizer and vllm-tgis-adapter and builds uv cache -# for transitive dependencies - eg. grpcio -############################################################### - -FROM base-builder AS vllmcache-builder - -ENV LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config -ENV PATH=/usr/lib64/llvm15/bin:$PATH - -COPY --from=torch-builder /tmp/control /dev/null -COPY --from=arrow-builder /tmp/control /dev/null -COPY --from=cv-builder /tmp/control /dev/null -COPY --from=numa-builder /tmp/control /dev/null -COPY --from=numba-builder /tmp/control /dev/null - ARG VLLM_TARGET_DEVICE=cpu -ARG GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 -# this step installs vllm and populates uv cache -# with all the transitive dependencies +USER root +WORKDIR /root + +ENV HOME=/root \ + WHEEL_DIR=/wheelsdir \ + VIRTUAL_ENV=/opt/vllm \ + GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 \ + CARGO_HOME=/root/.cargo \ + RUSTUP_HOME=/root/.rustup \ + UV_CACHE_DIR=$HOME/.cache/uv \ + PATH=/root/.cargo/bin:/root/.rustup/bin:${VIRTUAL_ENV}/bin:$PATH + +RUN echo "DEBUG: VLLM_VERSION=${VLLM_VERSION}" +RUN --mount=type=cache,target=/var/cache/dnf \ + microdnf install -y \ + python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip \ + && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ + && python${PYTHON_VERSION} -m pip install -U pip uv --no-cache + +# Important: Copy only bare minimum required for the script to run +COPY requirements/ requirements/ +COPY pyproject.toml ./ + +# The script is expected to install whatever python dependencies are missing +# as well as whatever system libraries need to be installed from source +COPY build_vllm_*.sh ./ + RUN --mount=type=cache,target=/root/.cache/uv \ - dnf install llvm15 llvm15-devel -y && \ - rpm -ivh --nodeps https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-16.el9.ppc64le.rpm && \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone https://github.com/huggingface/xet-core.git && cd xet-core/hf_xet/ && \ - uv pip install maturin && \ - uv build --wheel --out-dir /hf_wheels/ + sh ./build_vllm_$(uname -m).sh + +# copy vllm source code to build cache +COPY . . -ENV CXXFLAGS="-fno-lto -Wno-error=free-nonheap-object" \ - CFLAGS="-fno-lto" RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \ - --mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \ - --mount=type=bind,from=cv-builder,source=/opencvwheels/,target=/opencvwheels/,ro \ - --mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \ - --mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \ - --mount=type=bind,src=.,dst=/src/,rw \ source /opt/rh/gcc-toolset-14/enable && \ - export PATH=$PATH:/usr/lib64/llvm15/bin && \ - uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl && \ - sed -i -e 's/.*torch.*//g' /src/pyproject.toml /src/requirements/*.txt && \ - sed -i -e 's/.*sentencepiece.*//g' /src/pyproject.toml /src/requirements/*.txt && \ - uv pip install sentencepiece==0.2.0 pandas pythran nanobind pybind11 /hf_wheels/*.whl && \ - 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/cuda.txt --no-build-isolation && \ - cd /src/ && \ - uv build --wheel --out-dir /vllmwheel/ --no-build-isolation && \ - uv pip install /vllmwheel/*.whl + pip install -U uv +# build & install vLLM so that all transitive dependencies are build/downloaded into the uv cache +RUN --mount=type=cache,target=/root/.cache/uv \ + source /opt/rh/gcc-toolset-14/enable && \ + export PATH=/opt/rh/gcc-toolset-14/root/usr/bin:$PATH && \ + export CC=/opt/rh/gcc-toolset-14/root/usr/bin/gcc && \ + export CXX=/opt/rh/gcc-toolset-14/root/usr/bin/g++ && \ + export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/lib64/pkgconfig:$PKG_CONFIG_PATH && \ + export CMAKE_PREFIX_PATH=/usr/local:/usr:$CMAKE_PREFIX_PATH && \ + export Protobuf_PROTOC_EXECUTABLE=/usr/bin/protoc && \ + export CFLAGS="-mcpu=power10 -mtune=power10" && \ + export CXXFLAGS="-mcpu=power10 -mtune=power10" && \ + export DNNL_ARCH_OPT_FLAGS="-mcpu=power10 -mtune=power10" && \ + export C_INCLUDE_PATH=/usr/local/include:$C_INCLUDE_PATH && \ + export CPLUS_INCLUDE_PATH=/usr/local/include:$CPLUS_INCLUDE_PATH && \ + uv pip install 'setuptools>=78.1.1' && \ + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/OpenBLAS/lib/:/usr/local/lib64:/usr/local/lib && \ + export LIBGOMP=/opt/rh/gcc-toolset-14/root/usr/lib/gcc/ppc64le-redhat-linux/14/libgomp.so && \ -############################################################### -# Stage to build lapack -############################################################### - -FROM base-builder AS lapack-builder - -ARG MAX_JOBS -ARG LAPACK_VERSION=3.12.1 -RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${LAPACK_VERSION} \ - && cd lapack && source /opt/rh/gcc-toolset-14/enable \ - && cmake -B build -S . \ - && cmake --build build -j ${MAX_JOBS:-$(nproc)} + export CMAKE_LIBRARY_PATH=$(dirname $LIBGOMP):${CMAKE_LIBRARY_PATH} && \ + export LIBRARY_PATH=$(dirname $LIBGOMP):${LIBRARY_PATH} && \ + export LD_LIBRARY_PATH=$(dirname $LIBGOMP):${LD_LIBRARY_PATH} && \ + echo "LIBGOMP=${LIBGOMP}" && \ + find /root/.cache/uv -name "*.whl" && \ + SETUPTOOLS_SCM_PRETEND_VERSION="$VLLM_VERSION" uv build \ + --wheel --out-dir ${WHEEL_DIR} --no-build-isolation && \ + uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh ############################################################### # FINAL VLLM IMAGE STAGE # @@ -278,72 +83,74 @@ RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${L FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS vllm-openai ARG PYTHON_VERSION=3.12 -ARG OPENBLAS_VERSION=0.3.30 +ENV VLLM_NO_USAGE_STATS=1 # Set Environment Variables for venv & openblas ENV VIRTUAL_ENV=/opt/vllm -ENV PATH=${VIRTUAL_ENV}/bin:$PATH -ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig/ -ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib64:/usr/local/lib:/usr/lib64:/usr/lib +ENV PCP_DIR=/opt/rh/gcc-toolset-14/root +ENV PATH=${VIRTUAL_ENV}/bin:${PCP_DIR}/usr/bin:/usr/local/bin:$PATH +ENV PKG_CONFIG_PATH=${PCP_DIR}/usr/lib64/pkgconfig:/usr/local/lib/pkgconfig/ +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" +ENV LD_LIBRARY_PATH=${PCP_DIR}/usr/lib64:${PCP_DIR}/usr/lib:${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/torch/lib:/usr/local/lib:$LD_LIBRARY_PATH:/usr/local/lib64:/usr/lib64:/usr/lib ENV UV_LINK_MODE=copy -ENV OMP_NUM_THREADS=16 +ARG VLLM_VERSION="0.22.1" +ARG UV_EXTRA_INDEX_URL="https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/" +ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL} +ENV UV_INDEX_STRATEGY=first-match -# create artificial dependencies between stages for independent stages to build in parallel -COPY --from=torch-builder /tmp/control /dev/null -COPY --from=arrow-builder /tmp/control /dev/null -COPY --from=cv-builder /tmp/control /dev/null -COPY --from=vllmcache-builder /tmp/control /dev/null -COPY --from=numa-builder /tmp/control /dev/null -COPY --from=lapack-builder /tmp/control /dev/null -COPY --from=openblas-builder /tmp/control /dev/null -COPY --from=numba-builder /tmp/control /dev/null -# install gcc-11, python, openblas, numactl, lapack RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \ - --mount=type=bind,from=lapack-builder,source=/lapack/,target=/lapack/,rw \ - --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \ rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ microdnf install --nodocs -y \ - libomp libicu tar findutils openssl llvm15 llvm15-devel \ - pkgconfig xsimd g++ gcc-fortran libsndfile \ + libomp libicu tar autoconf automake libtool findutils openssl numactl numactl-devel \ + pkgconfig xsimd gcc-toolset-14 libsndfile \ libtiff libjpeg openjpeg2 zlib zeromq \ freetype lcms2 libwebp tcl tk utf8proc \ - harfbuzz fribidi libraqm libimagequant libxcb util-linux \ + harfbuzz fribidi libraqm libimagequant libxcb util-linux gperftools-libs \ python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip \ - && export PATH=$PATH:/usr/lib64/llvm15/bin && microdnf clean all \ - && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ - && python -m pip install -U pip uv --no-cache \ - && make -C /numactl install \ - && PREFIX=/usr/local make -C /openblas install \ - && uv pip install 'cmake<4' \ - && cmake --install /lapack/build \ - && uv pip uninstall cmake + && source /opt/rh/gcc-toolset-14/enable \ + && microdnf update -y \ + && microdnf clean all -# consume previously built wheels (including vllm) -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \ - --mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \ - --mount=type=bind,from=cv-builder,source=/opencvwheels/,target=/opencvwheels/,ro \ - --mount=type=bind,from=vllmcache-builder,source=/hf_wheels/,target=/hf_wheels/,ro \ - --mount=type=bind,from=vllmcache-builder,source=/vllmwheel/,target=/vllmwheel/,ro \ - --mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \ - export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':') && uv pip install sentencepiece==0.2.0 && \ - HOME=/root uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl /hf_wheels/*.whl /vllmwheel/*.whl +# The `lscpu` command was added as a requirement in part of https://github.com/vllm-project/vllm/pull/21032, so installing it. +RUN microdnf install --nodocs -y util-linux && \ + microdnf clean all +COPY --from=builder-base /usr/lib64/libprotobuf.so.25 /usr/lib64/ +COPY --from=builder-base /usr/lib64/libprotobuf.so.25.0.0 /usr/lib64/ -COPY ./ /workspace/vllm -WORKDIR /workspace/vllm -ARG GIT_REPO_CHECK=0 -RUN --mount=type=bind,source=.git,target=.git \ - if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi +# Use builder venv in final stage instead of wheel reinstallation +COPY --from=builder-base /opt/vllm /opt/vllm -# install development dependencies (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -e tests/vllm_test_utils +ENV LD_PRELOAD=/usr/lib64/libtcmalloc.so.4 -WORKDIR /workspace/ +WORKDIR /home/vllm -RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks +# setup non-root user for OpenShift +RUN umask 002 && \ + useradd --uid 2000 --gid 0 vllm && \ + mkdir -p /home/vllm && \ + chmod g+rwx /home/vllm + +ENV HOME=/home/vllm + +# Add labels to document build configuration +LABEL org.opencontainers.image.title="vLLM CPU" +LABEL org.opencontainers.image.description="vLLM inference engine for CPU platforms" +LABEL org.opencontainers.image.vendor="vLLM Project" +LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" + +# Build configuration labels +ARG TARGETARCH +ARG VLLM_CPU_PPC64LE +ARG PYTHON_VERSION + +LABEL ai.vllm.build.target-arch="${TARGETARCH}" +LABEL ai.vllm.build.cpu-ppc64le="${VLLM_CPU_PPC64LE:-false}" +LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}" + +USER 2000 ENTRYPOINT ["vllm", "serve"] + + diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 16284e999a7..93daf9d150d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,5 +1,7 @@ # default base image ARG REMOTE_VLLM="0" +ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" +ARG VLLM_BRANCH="main" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base @@ -110,10 +112,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" -ENV VLLM_REPO=${VLLM_REPO} -ENV VLLM_BRANCH=${VLLM_BRANCH} +ARG VLLM_REPO +ARG VLLM_BRANCH ONBUILD RUN git clone ${VLLM_REPO} \ && cd vllm \ && git fetch -v --prune -- origin ${VLLM_BRANCH} \ @@ -124,24 +124,96 @@ ONBUILD RUN git clone ${VLLM_REPO} \ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # ----------------------- -# Rust build stage -# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain or protoc. -FROM fetch_vllm AS rust-build -ARG COMMON_WORKDIR +# Rust/protoc toolchain inputs +# +# Keep these separate from Rust source inputs so Rust source-only changes do not +# invalidate the long-lived ci_base image that carries the toolchain. +FROM scratch AS rust_toolchain_input_0 +COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml +COPY tools/install_protoc.sh /rust-toolchain-input/vllm/tools/install_protoc.sh + +FROM base AS rust_toolchain_input_1 +ARG VLLM_REPO +ARG VLLM_BRANCH +RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \ + && cd /rust-toolchain-input/vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git sparse-checkout init --no-cone \ + && git sparse-checkout set \ + rust-toolchain.toml \ + tools/install_protoc.sh \ + && git checkout FETCH_HEAD + +FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input + +# ----------------------- +# Rust build inputs +# +# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts. +# REMOTE_VLLM=1: fetch only those same files from the requested remote ref. +FROM scratch AS rust_input_0 +COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt +COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/ +COPY rust/proto /rust-input/vllm/rust/proto +COPY rust/src /rust-input/vllm/rust/src +COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml +COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py +COPY build_rust.sh /rust-input/vllm/build_rust.sh + +FROM base AS rust_input_1 +ARG VLLM_REPO +ARG VLLM_BRANCH +RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \ + && cd /rust-input/vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git sparse-checkout init --no-cone \ + && git sparse-checkout set \ + requirements/build/rust.txt \ + rust/Cargo.lock \ + rust/Cargo.toml \ + rust/proto/** \ + rust/src/** \ + rust-toolchain.toml \ + tools/build_rust.py \ + build_rust.sh \ + && git checkout FETCH_HEAD + +FROM rust_input_${REMOTE_VLLM} AS rust-input + +# ----------------------- +# Rust/protoc toolchain +FROM base AS rust-toolchain + +ENV CARGO_HOME=/root/.cargo +ENV RUSTUP_HOME=/root/.rustup +ENV PATH=${CARGO_HOME}/bin:${PATH} # protoc is used by tonic-build/prost-build. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ ca-certificates curl unzip \ && rm -rf /var/lib/apt/lists/* -COPY tools/install_protoc.sh /tmp/install_protoc.sh +COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" +COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml +RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \ + && if ! command -v rustup >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --default-toolchain none; \ + fi \ + && rustup toolchain install "${TOOLCHAIN}" \ + && rustup default "${TOOLCHAIN}" \ + && rustc --version \ + && cargo --version + +# ----------------------- +# Rust build stage +# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages +# don't need the rust toolchain or protoc. +FROM rust-toolchain AS rust-build +ARG COMMON_WORKDIR +ARG USE_SCCACHE # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). @@ -149,16 +221,37 @@ ENV CARGO_BUILD_JOBS=4 ENV CARGO_NET_RETRY=10 ENV RUSTUP_MAX_RETRIES=10 +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt + +# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on +# a separate local sccache daemon while sharing the same remote cache backend. +ENV SCCACHE_SERVER_PORT=4227 + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/build/rust.txt + +COPY --from=rust-input /rust-input/vllm/rust rust +COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml +COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py +COPY --from=rust-input /rust-input/vllm/build_rust.sh build_rust.sh + # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while -# keeping the cache benefit. Copy the binary out so it persists into the image -# layer for later COPY --from=rust-build. +# keeping the cache benefit. Do not cache target/, because stale target metadata +# can outlive source updates across BuildKit cache reuse. RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ - --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \ - cd ${COMMON_WORKDIR}/vllm \ - && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \ - && test -x /tmp/vllm-rs + if [ "$USE_SCCACHE" = "1" ]; then \ + export RUSTC_WRAPPER=sccache \ + && sccache --show-stats; \ + fi \ + && bash build_rust.sh \ + && test -x vllm/vllm-rs \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + sccache --show-stats; \ + fi # ----------------------- # vLLM native build stages @@ -175,9 +268,10 @@ COPY requirements/common.txt requirements/common.txt RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ uv pip install --system -r requirements/rocm.txt -# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do -# not invalidate the expensive native build layer. +# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image, +# but it still participates in the BuildKit cache key for this RUN step. COPY setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py COPY cmake cmake/ COPY csrc csrc/ COPY vllm/envs.py vllm/envs.py @@ -209,9 +303,10 @@ ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ @@ -231,9 +326,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -418,9 +521,10 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ # Create /install directory for custom wheels RUN mkdir -p /install @@ -515,9 +619,17 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmar COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -527,6 +639,17 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 FROM mori_base AS ci_base ARG COMMON_WORKDIR +ENV CARGO_HOME=/root/.cargo +ENV RUSTUP_HOME=/root/.rustup +ENV PATH=${CARGO_HOME}/bin:${PATH} + +# Bake the Rust/protoc toolchain into ci_base so Rust build/test steps do not +# need to fetch it on each per-commit image build. +COPY --from=rust-toolchain /root/.cargo /root/.cargo +COPY --from=rust-toolchain /root/.rustup /root/.rustup +COPY --from=rust-toolchain /usr/local/bin/protoc /usr/local/bin/protoc +COPY --from=rust-toolchain /usr/local/include/google /usr/local/include/google + # Update rdma-core to support latest rocshmem. ARG DEEPEP_NIC RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ @@ -551,6 +674,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + unzip \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* @@ -574,6 +698,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_XET_HIGH_PERFORMANCE=1 ENV HF_HUB_DOWNLOAD_TIMEOUT=60 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Pre-install vLLM test dependencies. COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt RUN --mount=type=cache,target=/root/.cache/uv \ @@ -694,6 +821,9 @@ ENV SAFETENSORS_FAST_GPU=1 # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 208ce863f6b..2faaf774cf6 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,7 +1,7 @@ ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="ba5c1517" +ARG TRITON_BRANCH="0f380657" ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17 +ARG PYTORCH_BRANCH="d0c8b1f3" # release/2.11 as of 6/09 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" ARG PYTORCH_VISION_BRANCH="v0.24.1" ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" @@ -9,7 +9,7 @@ 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.13.post1" +ARG AITER_BRANCH="v0.1.16.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -114,12 +114,10 @@ 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 @@ -246,7 +244,7 @@ RUN pip install pyyaml && cd aiter \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ && sccache --show-stats; \ fi \ - && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index 554a7257c23..c2adfafdee2 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -37,37 +37,6 @@ RUN microdnf install -y \ python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip python${PYTHON_VERSION}-wheel && \ python${PYTHON_VERSION} -m venv $VIRTUAL_ENV && pip install --no-cache -U pip wheel uv && microdnf clean all -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 -b maint-19.0.1 && \ - cd arrow/cpp && \ - mkdir release && cd release && \ - cmake -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DARROW_PYTHON=ON \ - -DARROW_PARQUET=ON \ - -DARROW_ORC=ON \ - -DARROW_FILESYSTEM=ON \ - -DARROW_WITH_LZ4=ON \ - -DARROW_WITH_ZSTD=ON \ - -DARROW_WITH_SNAPPY=ON \ - -DARROW_JSON=ON \ - -DARROW_CSV=ON \ - -DARROW_DATASET=ON \ - -DPROTOBUF_PROTOC_EXECUTABLE=/usr/bin/protoc \ - -DARROW_DEPENDENCY_SOURCE=BUNDLED \ - .. && \ - make -j$(nproc) && \ - make install && \ - cd ../../python && \ - export PYARROW_PARALLEL=4 && \ - export ARROW_BUILD_TYPE=release && \ - 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 rust ENV CARGO_HOME=/root/.cargo ENV RUSTUP_HOME=/root/.rustup @@ -121,60 +90,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \ FROM python-install AS numba-builder ARG MAX_JOBS -ARG NUMBA_VERSION=0.61.2 +ARG NUMBA_VERSION=0.65.0 WORKDIR /tmp # Clone all required dependencies RUN --mount=type=cache,target=/root/.cache/uv \ microdnf install ninja-build gcc gcc-c++ -y && \ - git clone --recursive https://github.com/llvm/llvm-project.git -b llvmorg-15.0.7 && \ - git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \ + git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \ git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ - cd llvm-project && mkdir build && cd build && \ + cd llvmlite && \ 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')" \ - CMAKE_ARGS="${CMAKE_ARGS} -DFFI_INCLUDE_DIR=$PREFIX/include" \ - CMAKE_ARGS="${CMAKE_ARGS} -DFFI_LIBRARY_DIR=$PREFIX/lib" \ - cmake -DCMAKE_INSTALL_PREFIX="${PREFIX}" \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_LIBRARY_PATH="${PREFIX}" \ - -DLLVM_ENABLE_LIBEDIT=OFF \ - -DLLVM_ENABLE_LIBXML2=OFF \ - -DLLVM_ENABLE_RTTI=ON \ - -DLLVM_ENABLE_TERMINFO=OFF \ - -DLLVM_INCLUDE_BENCHMARKS=OFF \ - -DLLVM_INCLUDE_DOCS=OFF \ - -DLLVM_INCLUDE_EXAMPLES=OFF \ - -DLLVM_INCLUDE_GO_TESTS=OFF \ - -DLLVM_INCLUDE_TESTS=OFF \ - -DLLVM_INCLUDE_UTILS=ON \ - -DLLVM_INSTALL_UTILS=ON \ - -DLLVM_UTILS_INSTALL_DIR=libexec/llvm \ - -DLLVM_BUILD_LLVM_DYLIB=OFF \ - -DLLVM_LINK_LLVM_DYLIB=OFF \ - -DLLVM_EXPERIMENTAL_TARGETS_TO_BUILD=WebAssembly \ - -DLLVM_ENABLE_FFI=ON \ - -DLLVM_ENABLE_Z3_SOLVER=OFF \ - -DLLVM_OPTIMIZED_TABLEGEN=ON \ - -DCMAKE_POLICY_DEFAULT_CMP0111=NEW \ - -DCOMPILER_RT_BUILD_BUILTINS=ON \ - -DCOMPILER_RT_BUILTINS_HIDE_SYMBOLS=OFF \ - -DCOMPILER_RT_BUILD_LIBFUZZER=OFF \ - -DCOMPILER_RT_BUILD_CRT=OFF \ - -DCOMPILER_RT_BUILD_MEMPROF=OFF \ - -DCOMPILER_RT_BUILD_PROFILE=OFF \ - -DCOMPILER_RT_BUILD_SANITIZERS=OFF \ - -DCOMPILER_RT_BUILD_XRAY=OFF \ - -DCOMPILER_RT_BUILD_GWP_ASAN=OFF \ - -DCOMPILER_RT_BUILD_ORC=OFF \ - -DCOMPILER_RT_INCLUDE_TESTS=OFF \ - ${CMAKE_ARGS} -GNinja ../llvm \ - && ninja install . && \ - # build llvmlite - cd ../../llvmlite && python setup.py bdist_wheel && \ + python setup.py bdist_wheel && \ cd ../numba && \ if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \ sed -i '/#include "internal\/pycore_atomic.h"/i\#include "dynamic_annotations.h"' numba/_dispatcher.cpp; \ @@ -192,21 +119,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ cd opencv-python && \ python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist -## 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" -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 \ - 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 --compatibility linux - # # Final build stage FROM python-install AS vllm-cpu ARG PYTHON_VERSION @@ -235,28 +147,24 @@ RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \ 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 \ - --mount=type=bind,from=pyarrow,source=/tmp/arrow/python/dist,target=/tmp/arrow-wheels \ --mount=type=bind,from=torch-vision,source=/tmp/vision/dist,target=/tmp/vision-wheels/ \ --mount=type=bind,from=hf-xet-builder,source=/tmp/hf-xet/dist,target=/tmp/hf-xet-wheels/ \ --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=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) && \ - GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \ - uv pip install -v \ + uv pip install -v \ $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ $LLVM_WHL_FILE \ $NUMBA_WHL_FILE \ $OPENCV_WHL_FILE \ - $GUIDANCE_WHL_FILE \ + --torch-backend cpu \ --index-strategy unsafe-best-match \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 7b19e20f547..3bd16e8629b 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -6,44 +6,48 @@ ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update -y \ && apt-get install -y --no-install-recommends \ - ca-certificates curl git build-essential unzip \ + ca-certificates curl git build-essential unzip python3 python3-pip \ && rm -rf /var/lib/apt/lists/* COPY tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" - WORKDIR /workspace +COPY requirements/build/rust.txt requirements/build/rust.txt +RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt + +# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed +# by the wheel build stage. COPY rust rust COPY rust-toolchain.toml rust-toolchain.toml +COPY tools/build_rust.py tools/build_rust.py COPY build_rust.sh build_rust.sh # Cap cargo parallelism to avoid exhausting the CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ - --mount=type=cache,target=/workspace/rust/target \ - VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh +RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + bash build_rust.sh -FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base +FROM ubuntu:24.04 AS vllm-base + +ENV DEBIAN_FRONTEND=noninteractive WORKDIR /workspace/ ARG PYTHON_VERSION=3.12 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" -RUN apt clean && apt-get update -y && \ - apt-get install -y --no-install-recommends --fix-missing \ +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ curl \ ffmpeg \ git \ + gpg \ libsndfile1 \ libsm6 \ libxext6 \ @@ -53,9 +57,11 @@ RUN apt clean && apt-get update -y && \ numactl \ wget \ vim \ + ca-certificates \ python3.12 \ python3.12-dev \ - python3-pip + python3-pip && \ + rm -rf /var/lib/apt/lists/* # Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer. RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ @@ -75,13 +81,14 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO # Install UMD RUN mkdir neo && \ cd neo && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \ + wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \ + wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \ + wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \ + wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero-devel_1.28.2+u24.04_amd64.deb && \ dpkg -i *.deb && \ cd .. && \ rm -rf neo @@ -101,7 +108,13 @@ RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \ rm -f /opt/intel/oneapi/ccl/latest && \ - ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest + ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \ + printf '%s\n' \ + '/opt/intel/oneapi/ccl/2021.15/lib' \ + '/opt/intel/oneapi/mpi/2021.15/lib' \ + '/opt/intel/oneapi/compiler/2025.3/lib' \ + > /etc/ld.so.conf.d/oneapi-ccl.conf && \ + ldconfig SHELL ["bash", "-c"] CMD ["bash", "-c", "source /root/.bashrc && exec bash"] @@ -123,7 +136,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ -ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/" +ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib CMD ["/bin/bash"] ######################### UCX + NIXL BUILD STAGE ######################### @@ -131,8 +144,8 @@ CMD ["/bin/bash"] # never included in the final runtime image (mirrors ROCm's build_rixl stage). FROM vllm-base AS ucx-nixl-build -ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349 -ARG NIXL_VERSION=0.7.0 +ARG UCX_VERSION=v1.21.0-rc2 +ARG NIXL_VERSION=v1.2.0 # Build-time only: compiler, autotools, and verbs dev headers RUN apt-get update -y && apt-get install -y --no-install-recommends \ @@ -149,26 +162,25 @@ RUN apt-get update -y && apt-get install -y --no-install-recommends \ # patchelf (installed via uv) is used by the NIXL wheel build to rewrite # RPATH entries, making the wheel portable across stages. RUN --mount=type=cache,target=/root/.cache/uv \ - git clone https://github.com/openucx/ucx /tmp/ucx_source && \ - cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \ + git clone --depth 1 --branch "${UCX_VERSION}" https://github.com/openucx/ucx /tmp/ucx_source && \ + cd /tmp/ucx_source && \ bash autogen.sh && \ ./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \ - make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \ - git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \ - cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \ + make CFLAGS="-Wno-error=incompatible-pointer-types" -j"$(nproc)" && make install && \ + git clone --depth 1 --branch "${NIXL_VERSION}" https://github.com/ai-dynamo/nixl /tmp/nixl_source && \ + cd /tmp/nixl_source && \ uv pip install --upgrade meson pybind11 patchelf && \ uv pip install -r requirements.txt && \ PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \ LD_LIBRARY_PATH=/tmp/ucx_install/lib \ python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \ find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \ - rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \ - rm -rf /tmp/ucx_source /tmp/nixl_source + rm -rf /tmp/ucx_install/{include,share,etc,bin} /tmp/ucx_install/lib/cmake \ + /tmp/ucx_source /tmp/nixl_source FROM vllm-base AS vllm-openai -ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349 -ARG NIXL_VERSION=0.7.0 +ARG NIXL_VERSION=v1.2.0 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel. # No compiler or autotools are installed in this stage. @@ -192,7 +204,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ibverbs-providers \ librdmacm1t64 \ && rm -rf /var/lib/apt/lists/* \ - && uv pip install --no-deps /tmp/nixl_wheels/nixl-*.whl \ + && uv pip install --no-deps /tmp/nixl_wheels/nixl*.whl \ + && uv pip install nixl==${NIXL_VERSION} && uv pip uninstall nixl-cu13 \ && rm -rf /tmp/nixl_wheels RUN --mount=type=cache,target=/root/.cache/uv \ @@ -204,18 +217,16 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \ uv pip uninstall triton triton-xpu && \ uv pip install triton-xpu==3.7.1 && \ - uv pip uninstall oneccl oneccl-devel && \ - 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 uninstall oneccl oneccl-devel # Keep source-dependent layers near the end so frequent code-only changes # don't invalidate heavy dependency and UCX/NIXL layers. COPY . . -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index 138adcffcad..0ae991bd52d 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -141,6 +141,10 @@ variable "ROCM_CSRC_CACHE_TO_MODE" { default = "max" } +variable "ROCM_RUST_CACHE_TO_MODE" { + default = "max" +} + variable "ROCM_FINAL_CACHE_TO_MODE" { default = "min" } @@ -164,6 +168,13 @@ function "get_cache_from_rocm" { VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", + # Import the source-scoped Rust frontend cache so non-Rust changes do not + # force a fresh cargo release build. + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "", + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "", + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", # Branch-scoped full image cache - fallback when parent-commit cache is evicted ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", @@ -204,6 +215,27 @@ function "get_cache_to_rocm_csrc" { ]) } +function "get_cache_from_rocm_rust" { + params = [] + result = compact([ + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "", + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "", + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", + ]) +} + +function "get_cache_to_rocm_rust" { + params = [] + result = compact([ + # Export exact-commit and branch-scoped Rust caches. A content-addressed + # cache ref is appended by ci-bake-rocm.sh when that wrapper is used. + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT},mode=${ROCM_RUST_CACHE_TO_MODE}" : "", + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_RUST_CACHE_TO_MODE}" : "", + ]) +} + # Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP). # These stages are pinned to specific upstream commit hashes, so cache keys use # those hashes rather than the Buildkite commit. This means the cache persists @@ -276,6 +308,17 @@ target "csrc-rocm-ci" { output = ["type=cacheonly"] } +# Cache-only target for the Rust frontend build stage. Final-image cache +# exports use mode=min and do not reliably persist intermediate cargo layers, +# so Rust gets its own source-scoped cache target. +target "rust-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "rust-build" + cache-from = get_cache_from_rocm_rust() + cache-to = get_cache_to_rocm_rust() + output = ["type=cacheonly"] +} + # Keep wheel export on the same CI graph as the test image build so the # shared build_vllm/export_vllm stages resolve identically within one bake # invocation. Without this, export-wheel-rocm uses the plain local target @@ -292,24 +335,32 @@ target "export-wheel-rocm" { # Artifact-only vLLM build. GPU test jobs consume this artifact on top of # ci_base, avoiding a per-commit multi-GB image push/pull. group "test-rocm-ci-with-artifacts" { - targets = ["csrc-rocm-ci", "export-wheel-rocm"] + targets = ["rust-rocm-ci", "csrc-rocm-ci", "export-wheel-rocm"] } # Full test image + wheel export. Kept for fallback/debugging when a pushed # per-commit image is useful. group "test-rocm-ci-with-wheel" { - targets = ["csrc-rocm-ci", "test-rocm-ci", "export-wheel-rocm"] + targets = ["rust-rocm-ci", "csrc-rocm-ci", "test-rocm-ci", "export-wheel-rocm"] } # Image tags for the ci_base build. ci-bake-rocm.sh rewrites CI_BASE_IMAGE_TAG -# to the primary tag for this build. Non-nightly builds use a commit-scoped tag -# and also publish a content tag for reuse. NIGHTLY=1 builds on the stable branch -# can additionally set CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base. +# to the primary tag for this build. Builds always publish a content-scoped tag +# when the ci_base content hash is available. Builds with BUILDKITE_COMMIT also +# publish a commit-scoped tag, either as the primary tag or an additional alias. +# NIGHTLY=1 builds on the stable branch can additionally set +# CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base. variable "CI_BASE_IMAGE_TAG" { default = "rocm/vllm-dev:ci_base" } -variable "CI_BASE_IMAGE_TAG_CONTENT" { +# Supplemental tags only. ci-bake-rocm.sh leaves these empty when the same ref +# is already the primary CI_BASE_IMAGE_TAG. +variable "CI_BASE_IMAGE_TAG_COMMIT_EXTRA" { + default = "" +} + +variable "CI_BASE_IMAGE_TAG_CONTENT_EXTRA" { default = "" } @@ -357,7 +408,8 @@ target "ci-base-rocm-ci" { cache-from = concat( compact([ CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "", - CI_BASE_IMAGE_TAG_CONTENT != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT}" : "", + CI_BASE_IMAGE_TAG_COMMIT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_COMMIT_EXTRA}" : "", + CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "", CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "", ]), # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages @@ -365,7 +417,7 @@ target "ci-base-rocm-ci" { get_cache_from_rocm_deps(), ) cache-to = ["type=inline"] - tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE]) + tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_COMMIT_EXTRA, CI_BASE_IMAGE_TAG_CONTENT_EXTRA, CI_BASE_IMAGE_TAG_STABLE]) output = ["type=registry"] } diff --git a/docker/versions.json b/docker/versions.json index 15f77648a9c..4dffa00985c 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -34,15 +34,6 @@ "INSTALL_KV_CONNECTORS": { "default": "false" }, - "TORCH_CUDA_ARCH_LIST": { - "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX" - }, - "MAX_JOBS": { - "default": "2" - }, - "NVCC_THREADS": { - "default": "8" - }, "SCCACHE_BUCKET_NAME": { "default": "vllm-build-sccache" }, @@ -52,6 +43,15 @@ "SCCACHE_S3_NO_CREDENTIALS": { "default": "0" }, + "TORCH_CUDA_ARCH_LIST": { + "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0" + }, + "MAX_JOBS": { + "default": "2" + }, + "NVCC_THREADS": { + "default": "8" + }, "vllm_target_device": { "default": "cuda" }, @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.12" + "default": "0.6.13" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index b4f505493ad..8cb98a8f4e4 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/benchmarking/cli.md b/docs/benchmarking/cli.md index 6d0b2a01aca..692cb4918ba 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -37,7 +37,7 @@ th { | HuggingFace-HumanEval | ✅ | ✅ | `openai/openai_humaneval` | | HuggingFace-GSM8K | ✅ | ✅ | `openai/gsm8k` | | 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` | +| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` | | 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` | @@ -338,7 +338,7 @@ vllm bench serve \ --model meta-llama/Meta-Llama-3-8B-Instruct \ --dataset-name spec_bench \ --dataset-path "/data/spec_bench/question.jsonl" \ - --num-prompts -1 + --num-prompts -1 \ --spec-bench-category "summarization" ``` @@ -352,7 +352,7 @@ vllm bench serve \ 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 - +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: @@ -388,7 +388,7 @@ vllm bench serve \ --model meta-llama/Llama-3.3-70B-Instruct \ --dataset-name speed_bench \ --dataset-path "/data/speed_bench" \ - --num-prompts -1 + --num-prompts -1 \ --speed-bench-category "multilingual" ``` @@ -398,13 +398,54 @@ Run all categories in the Throughput split (2k ISL): vllm bench serve \ --model meta-llama/Llama-3.3-70B-Instruct \ --dataset-name speed_bench \ - --speed-bench-dataset-subset throughput_2k + --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. +#### BFCL (Tool-Calling) Benchmark + +The Berkeley Function Calling Leaderboard (BFCL) dataset measures serving +latency and throughput on realistic tool-calling traffic. Each request +carries a per-sample `tools` schema and chat history, so the server must +expose `/v1/chat/completions` with an auto-tool-choice parser enabled. +The benchmark client always uses the `openai-chat` backend. + +Start a tool-parser-enabled server, then run the bench. For example, with +`gpt-oss-20b`: + +```bash +# Server +vllm serve openai/gpt-oss-20b \ + --enable-auto-tool-choice \ + --tool-call-parser openai \ + --reasoning-parser openai_gptoss + +# Client +vllm bench serve \ + --backend openai-chat \ + --endpoint /v1/chat/completions \ + --model openai/gpt-oss-20b \ + --dataset-name hf \ + --dataset-path gorilla-llm/Berkeley-Function-Calling-Leaderboard \ + --bfcl-categories simple,live_simple,multiple \ + --num-prompts 200 +``` + +`--bfcl-categories` is a comma-separated list of BFCL v3 category names +(without the `BFCL_v3_` prefix or `.json` suffix). Defaults to +`simple,live_simple,multiple`. Other supported non-multi-turn categories +include `parallel`, `live_parallel`, `parallel_multiple`, +`live_parallel_multiple`, `irrelevance`, `live_irrelevance`, +`live_relevance`, `java`, `javascript`, and `rest`. Multi-turn categories +are not yet supported. + +The dataset class normalizes BFCL's loose schema dialect (`dict` → +`object`, `float` → `number`, `tuple` → `array`, `any` → `string`) so +modern grammar backends accept the translated tool definitions. + #### Other HuggingFaceDataset Examples ```bash @@ -491,7 +532,7 @@ vllm bench serve \ --blazedit-max-distance 0.99 ``` -`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` +`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` ```bash vllm bench serve \ @@ -1296,7 +1337,7 @@ Serve and benchmark VLM2Vec: # Run this in another process vllm serve TIGER-Lab/VLM2Vec-Full --runner pooling \ --trust-remote-code \ - --chat-template examples/template_vlm2vec_phi3v.jinja + --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja # Run these one by one after the server is up # download dataset diff --git a/docs/cli/README.md b/docs/cli/README.md index 08e986a7463..123f3f109a9 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -50,6 +50,21 @@ vllm serve --help=max-num-seqs vllm serve --help=max ``` +!!! tip "Human-readable integer arguments" + Many integer arguments accept human-readable suffixes for convenience. For example: + + - `1k` = 1,000 (decimal kilo) + - `1K` = 1,024 (binary kibibyte) + - `1m` = 1,000,000 (decimal mega) + - `1M` = 1,048,576 (binary mebibyte) + - `1g` / `1G` = 1 billion / 1 gibibyte + - `1t` / `1T` = 1 trillion / 1 tebibyte + + Decimal suffixes (`k`, `m`, `g`, `t`) also accept floating point: `25.6k` = 25,600. + Binary suffixes (`K`, `M`, `G`, `T`) require integers: `32K` = 32,768. + + Supported arguments include: `--max-model-len`, `--max-num-batched-tokens`, `--max-num-scheduled-tokens`, `--kv-cache-memory-bytes`, `--safetensors-prefetch-block-size`. + See [vllm serve](./serve.md) for the full reference of all available arguments. ## launch @@ -80,6 +95,9 @@ vllm chat --url http://{vllm-serve-host}:{vllm-serve-port}/v1 # Quick chat with a single prompt vllm chat --quick "hi" + +# Print TTFT and throughput statistics after each response +vllm chat --stats ``` See [vllm chat](./chat.md) for the full reference of all available arguments. @@ -97,6 +115,9 @@ vllm complete --url http://{vllm-serve-host}:{vllm-serve-port}/v1 # Quick complete with a single prompt vllm complete --quick "The future of AI is" + +# Print TTFT and throughput statistics after each response +vllm complete --stats ``` See [vllm complete](./complete.md) for the full reference of all available arguments. diff --git a/docs/configuration/conserving_memory.md b/docs/configuration/conserving_memory.md index 2c098118dbb..96a903bc31d 100644 --- a/docs/configuration/conserving_memory.md +++ b/docs/configuration/conserving_memory.md @@ -42,7 +42,7 @@ and the maximum batch size (`max_num_seqs` option). ```python from vllm import LLM -llm = LLM(model="adept/fuyu-8b", max_model_len=2048, max_num_seqs=2) +llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct", max_model_len=2048, max_num_seqs=2) ``` ## Reduce CUDA Graphs diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index 80aec64ee5b..efa0f8b9046 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -16,6 +16,14 @@ vLLM provides 4 optimization levels (`-O0`, `-O1`, `-O2`, `-O3`) that allow user For more information, see the [optimization level documentation](../design/optimization_levels.md). +## Faster Startup + +Beyond the optimization levels, three mechanisms reduce time-to-first-token on repeated boots of the same (model, config, hardware) combination: + +- **Reuse the compile cache.** vLLM persists `torch.compile` artifacts under `VLLM_CACHE_ROOT` (default `~/.cache/vllm`), and the cache directory can be copied between machines or baked into a container image; see the [torch.compile design doc](../design/torch_compile.md). Set `VLLM_FORCE_AOT_LOAD=1` to fail loudly instead of silently recompiling when the cache misses (any change to the model, config, relevant `VLLM_*` environment variables, torch build, or GPU model invalidates it). +- **Skip memory profiling with `--kv-cache-memory`.** On startup, vLLM logs the exact `--kv-cache-memory` value that reproduces the current allocation. Passing it back on the next boot skips the memory-profiling measurement and the CUDA-graph memory estimation pass. Note that this has performance implications: the KV cache is sized to exactly the given value instead of being measured, so a conservative value caps batch concurrency (and therefore throughput), while an optimistic one fails at allocation time. The value is only valid on the same GPU with the same initial free memory; if a boot OOMs after hardware or co-tenant changes, remove the flag to re-profile. +- **Serve without CUDA graphs using `--enforce-eager`.** Skips both compilation and CUDA-graph capture for the fastest possible startup, at the cost of steady-state decode performance. Useful for development loops and for measuring how much of a boot is compile/capture. + ## Preemption Due to the autoregressive nature of transformer architecture, there are times when KV cache space is insufficient to handle all batched requests. @@ -109,7 +117,7 @@ from vllm import LLM # Combine pipeline and tensor parallelism llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct, + model="meta-llama/Llama-3.3-70B-Instruct", tensor_parallel_size=4, pipeline_parallel_size=2, ) @@ -276,8 +284,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.) you can switch to the [fastokens](https://github.com/crusoecloud/fastokens) Rust backend, a drop-in replacement that's substantially faster on -encode/decode and on streaming detokenization. Enable it by setting -`VLLM_USE_FASTOKENS=1`: +encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is +available in vLLM v0.23.0 and later. If your installed vLLM version does not +recognize the environment variable, upgrade vLLM before enabling the override: ```console VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B @@ -296,8 +305,8 @@ llm = LLM(model="Qwen/Qwen3-8B") The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't, vLLM raises a clear `ImportError` at tokenizer load. The override applies to any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`, -`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Models that don't use the HF -fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag. +`deepseek_v32`, `deepseek_v4`, …). Models that don't use the HF +fast tokenizer (`mistral`, `kimi_audio`) ignore the flag. Tokenizer-bound workloads — long shared prefixes, bursty short prompts, batch detokenization — see the largest wins. If your bottleneck is GPU diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 3fc8b6dd52b..34dc385db78 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -304,9 +304,15 @@ review process: resources. The reviewer will add `ready` label to the PR when the PR is ready to merge or a full CI run is needed. -### Escalating Stalled Contributions +### Pull Request Limits and Escalation -If you have an important contribution that has not yet received maintainer attention, please email us at: +vLLM uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/) +for contributors without write access. The current cap is 6 open PRs. If this +blocks well-intentioned critical work, contact a committer to request bypass +list access. + +If you need an expedited review for an important contribution, please email us +at: diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index a0038f461a0..c57c430478f 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -60,15 +60,21 @@ the failure? ## Logs Wrangling -Download a job's log (no Buildkite login required): - +Logs are public; no Buildkite login needed. [.buildkite/scripts/ci-fetch-log.sh](../../../.buildkite/scripts/ci-fetch-log.sh) +saves each log as `ci--.log`, stripped of timestamps and +ANSI codes: ```bash -# Find the failing job. Each row's URL is .../builds/#: -gh pr checks --repo vllm-project/vllm +# All failed jobs in a PR's latest build (current branch's PR if omitted): +.buildkite/scripts/ci-fetch-log.sh --pr -# Download + strip timestamps/ANSI in one step: +# All failed jobs in a build (--soft also includes soft-failed jobs; +# --all fetches every finished job): +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/" + +# One job — `gh pr checks` URLs (#) and web UI URLs (?sid=) both +# work; pass "-" as a second argument to stream to stdout: .buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/#" ``` diff --git a/docs/contributing/ci/nightly_builds.md b/docs/contributing/ci/nightly_builds.md index 8f3512db3d4..10c4a437240 100644 --- a/docs/contributing/ci/nightly_builds.md +++ b/docs/contributing/ci/nightly_builds.md @@ -14,7 +14,7 @@ Wheels are built in the `Release` pipeline (`.buildkite/release-pipeline.yaml`) Each build step: 1. Builds the wheel in a Docker container. -2. Renames the wheel filename to use the correct manylinux tag (currently `manylinux_2_31`) for PEP 600 compliance. +2. Renames the wheel filename to use the correct manylinux tag (currently `manylinux_2_28`) for PEP 600 compliance. 3. Uploads the wheel to S3 bucket `vllm-wheels` under `/{commit_hash}/`. ### Index Generation diff --git a/docs/contributing/ci/update_pytorch_version.md b/docs/contributing/ci/update_pytorch_version.md index 98947dd4402..1c29bf2438c 100644 --- a/docs/contributing/ci/update_pytorch_version.md +++ b/docs/contributing/ci/update_pytorch_version.md @@ -93,6 +93,28 @@ To address this, manually trigger a build on Buildkite to accomplish two objecti Buildkite new build popup

+You can also trigger this build from the command line with +[`.buildkite/scripts/trigger-ci-build.sh`](../../../.buildkite/scripts/trigger-ci-build.sh) +(dry-run by default; pass `--execute` to actually trigger it). + +## Test against PyTorch nightly + +The steps above test against a specific PyTorch RC/stable wheel pinned in the +requirements files. To instead build and run the CI suite against the latest +PyTorch **nightly** wheels, set the `TORCH_NIGHTLY=1` environment variable on +the build (or apply the `ready-torch-nightly` label to the PR). + +When `TORCH_NIGHTLY=1`, the base CI image is built against PyTorch nightly +(`image_build_torch_nightly.sh`, `PYTORCH_NIGHTLY=1`, CUDA 13.0) and tagged at +the normal image tag, so the entire existing pipeline runs on nightly torch -- +there is no separate pipeline section to trigger. Combine it with `RUN_ALL=1` +to run the full suite (the `ready-torch-nightly` label and +`trigger-ci-build.sh --torch-nightly` both set this for you). This is the +configuration to use for a scheduled "vLLM vs PyTorch nightly" run. + +Use `.buildkite/scripts/trigger-ci-build.sh --torch-nightly` to trigger it from +the command line. + ## Update all the different vLLM platforms Rather than attempting to update all vLLM platforms in a single pull request, it's more manageable diff --git a/docs/contributing/model/basic.md b/docs/contributing/model/basic.md index dceb78f5263..0cc24baae92 100644 --- a/docs/contributing/model/basic.md +++ b/docs/contributing/model/basic.md @@ -133,10 +133,10 @@ The model should inherit protocol `IsAttentionFree` and also implement class met For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes. The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized. -For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together). +For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together). These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol). -For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively. +For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`. Please follow the same guidelines as case (2) for implementing these models. We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention). For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`. @@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this. It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend. Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it. -Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this. +Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this. The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended. diff --git a/docs/contributing/model/multimodal.md b/docs/contributing/model/multimodal.md index 67cde8df987..33d89db75d3 100644 --- a/docs/contributing/model/multimodal.md +++ b/docs/contributing/model/multimodal.md @@ -324,154 +324,44 @@ Assuming that the memory usage increases with the number of tokens, the dummy in return image_token * num_images ``` -=== "No input placeholders: Fuyu" +=== "No input placeholders: PaliGemma" - Looking at the code of HF's `FuyuForCausalLM`: - - ??? code - - ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/modeling_fuyu.py#L311-L322 - if image_patches is not None and past_key_values is None: - patch_embeddings = [ - self.vision_embed_tokens(patch.to(self.vision_embed_tokens.weight.dtype)) - .squeeze(0) - .to(inputs_embeds.device) - for patch in image_patches - ] - inputs_embeds = self.gather_continuous_embeddings( - word_embeddings=inputs_embeds, - continuous_embeddings=patch_embeddings, - image_patch_input_indices=image_patches_indices, - ) - ``` - - The number of placeholder feature tokens for the `i`th item in the batch is `patch_embeddings[i].shape[0]`, - which is the same as `image_patches[i].shape[0]`, i.e. `num_total_patches`. - - Unlike LLaVA, Fuyu does not define the number of patches inside the modeling file. Where can we get more information? - Considering that the model input comes from the output of `FuyuProcessor`, let's **look at the preprocessing files**. - - The image outputs are obtained by calling `FuyuImageProcessor.preprocess` and then - `FuyuImageProcessor.preprocess_with_tokenizer_info` inside `FuyuProcessor`. - - In `FuyuImageProcessor.preprocess`, the images are resized and padded to the target `FuyuImageProcessor.size`, - returning the dimensions after resizing (but before padding) as metadata. - - ??? code - - ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/processing_fuyu.py#L541-L544 - image_encoding = self.image_processor.preprocess(images, **output_kwargs["images_kwargs"]) - batch_images = image_encoding["images"] - image_unpadded_heights = image_encoding["image_unpadded_heights"] - image_unpadded_widths = image_encoding["image_unpadded_widths"] - - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/image_processing_fuyu.py#L480-L - if do_resize: - batch_images = [ - [self.resize(image, size=size, input_data_format=input_data_format) for image in images] - for images in batch_images - ] - - image_sizes = [get_image_size(images[0], channel_dim=input_data_format) for images in batch_images] - image_unpadded_heights = [[image_size[0]] for image_size in image_sizes] - image_unpadded_widths = [[image_size[1]] for image_size in image_sizes] - - if do_pad: - batch_images = [ - [ - self.pad_image( - image, - size=size, - mode=padding_mode, - constant_values=padding_value, - input_data_format=input_data_format, - ) - for image in images - ] - for images in batch_images - ] - ``` - - In `FuyuImageProcessor.preprocess_with_tokenizer_info`, the images are split into patches based on this metadata: - - ??? code - - ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/processing_fuyu.py#L417-L425 - model_image_input = self.image_processor.preprocess_with_tokenizer_info( - image_input=tensor_batch_images, - image_present=image_present, - image_unpadded_h=image_unpadded_heights, - image_unpadded_w=image_unpadded_widths, - image_placeholder_id=image_placeholder_id, - image_newline_id=image_newline_id, - variable_sized=True, - ) - - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/image_processing_fuyu.py#L638-L658 - image_height, image_width = image.shape[1], image.shape[2] - if variable_sized: # variable_sized=True - new_h = min( - image_height, - math.ceil(image_unpadded_h[batch_index, subseq_index] / patch_height) * patch_height, - ) - new_w = min( - image_width, - math.ceil(image_unpadded_w[batch_index, subseq_index] / patch_width) * patch_width, - ) - image = image[:, :new_h, :new_w] - image_height, image_width = new_h, new_w - - num_patches = self.get_num_patches(image_height=image_height, image_width=image_width) - tensor_of_image_ids = torch.full( - [num_patches], image_placeholder_id, dtype=torch.int32, device=image_input.device - ) - patches = self.patchify_image(image=image.unsqueeze(0)).squeeze(0) - assert num_patches == patches.shape[0] - ``` - - The number of patches is in turn defined by `FuyuImageProcessor.get_num_patches`: - - ??? code - - ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/image_processing_fuyu.py#L552-L562 - patch_size = patch_size if patch_size is not None else self.patch_size - patch_height, patch_width = self.patch_size["height"], self.patch_size["width"] - - if image_height % patch_height != 0: - raise ValueError(f"{image_height=} must be divisible by {patch_height}") - if image_width % patch_width != 0: - raise ValueError(f"{image_width=} must be divisible by {patch_width}") - - num_patches_per_dim_h = image_height // patch_height - num_patches_per_dim_w = image_width // patch_width - num_patches = num_patches_per_dim_h * num_patches_per_dim_w - ``` - - These image patches correspond to placeholder tokens (`|SPEAKER|`). So, we just need to maximize the number of image patches. Since input images are first resized - to fit within `image_processor.size`, we can maximize the number of image patches by inputting an image with size equal to `image_processor.size`. - - ```python - def get_image_size_with_most_features(self) -> ImageSize: - image_processor = self.get_image_processor() - return ImageSize( - width=image_processor.size["width"], - height=image_processor.size["height"], - ) - ``` - - Fuyu does not expect image placeholders in the inputs to HF processor, so - the dummy prompt text is empty regardless of the number of images. + Unlike LLaVA, PaliGemma's HF processor does not expect image placeholder + tokens in the input prompt; the placeholder feature tokens are instead + inserted afterwards (see [Prompt updates](#prompt-updates)). So the dummy + prompt text is empty regardless of the number of images: ```python def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "" ``` - For the multimodal image profiling data, the logic is very similar to LLaVA: + PaliGemma resizes every image to a square of `vision_config.image_size`, so + the number of placeholder feature tokens per image is fixed at + `(image_size // patch_size) ** 2`. This is computed by the SigLIP vision + encoder that PaliGemma uses: + + ??? code + + ```python + # vllm/model_executor/models/siglip.py + class SiglipEncoderInfo(VisionEncoderInfo[SiglipVisionConfig]): + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + return self.get_patch_grid_length() ** 2 + + def get_patch_grid_length(self) -> int: + image_size, patch_size = self.get_image_size(), self.get_patch_size() + return image_size // patch_size + ``` + + Since the number of image tokens doesn't depend on the input image dimensions, + we can simply use a dummy image of the model's expected input size for the + multimodal profiling data: ??? code @@ -482,16 +372,18 @@ Assuming that the memory usage increases with the number of tokens, the dummy in mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions], ) -> MultiModalDataDict: - target_width, target_height = \ - self.info.get_image_size_with_most_features() + hf_config = self.info.get_hf_config() + vision_config = hf_config.vision_config + max_image_size = vision_config.image_size + num_images = mm_counts.get("image", 0) image_overrides = mm_options.get("image") return { "image": self._get_dummy_images( - width=target_width, - height=target_height, + width=max_image_size, + height=max_image_size, num_images=num_images, overrides=image_overrides, ) @@ -545,28 +437,15 @@ return a schema of the tensors outputted by the HF processor that are related to Our [actual code](../../../vllm/model_executor/models/llava.py) additionally supports pre-computed image embeddings, which can be passed to be model via the `image_embeds` argument. -=== "With postprocessing: Fuyu" +=== "With postprocessing: Mistral3" - The `image_patches` output of `FuyuImageProcessor.preprocess_with_tokenizer_info` concatenates - the patches from each image belonging to an item in the batch: + The `pixel_values` output of Mistral3's HF processor pads every image in the + batch to a common size, so that they can be stacked into a single tensor. - ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/image_processing_fuyu.py#L673-L679 - image_input_ids.append(tensor_of_image_ids) - image_patches.append(patches) - else: - image_input_ids.append(torch.tensor([], dtype=torch.int32, device=image_input.device)) - - batch_image_input_ids.append(image_input_ids) - batch_image_patches.append(image_patches) - ``` - - The shape of `image_patches` outputted by `FuyuImageProcessor` is therefore - `(1, num_images, num_patches, patch_width * patch_height * num_channels)`. - - In order to support the use of - [MultiModalFieldConfig.batched][vllm.multimodal.inputs.MultiModalFieldConfig.batched] - like in LLaVA, we remove the extra batch dimension by overriding + To use [MultiModalFieldConfig.batched][vllm.multimodal.inputs.MultiModalFieldConfig.batched] + like in LLaVA, each image's features must be independent of the others (which + is also required for prefix caching to work correctly). So, we un-pad each image + back to its own size by overriding [BaseMultiModalProcessor._call_hf_processor][vllm.multimodal.processing.BaseMultiModalProcessor._call_hf_processor]: ??? code @@ -586,33 +465,27 @@ return a schema of the tensors outputted by the HF processor that are related to tok_kwargs=tok_kwargs, ) - image_patches = processed_outputs.get("image_patches") - if image_patches is not None: - images = mm_data["images"] - assert isinstance(images, list) + pixel_values = processed_outputs.get("pixel_values") + if pixel_values is not None: + # Avoid padding since we need the output for each image to be + # independent of other images for the cache to work correctly + image_sizes = processed_outputs["image_sizes"] + assert len(pixel_values) == len(image_sizes) - # Original output: (1, num_images, Pn, Px * Py * C) - # New output: (num_images, Pn, Px * Py * C) - assert (isinstance(image_patches, list) - and len(image_patches) == 1) - assert (isinstance(image_patches[0], torch.Tensor) - and len(image_patches[0]) == len(images)) - - processed_outputs["image_patches"] = image_patches[0] + processed_outputs["pixel_values"] = [ + p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes) + ] return processed_outputs ``` - !!! note - Our [actual code](../../../vllm/model_executor/models/fuyu.py) has special handling - for text-only inputs to prevent unnecessary warnings from HF processor. - !!! note The `_call_hf_processor` method specifies both `mm_kwargs` and `tok_kwargs` for processing. `mm_kwargs` is used to both initialize and call the huggingface processor, whereas `tok_kwargs` is only used to call the huggingface processor. - This lets us override [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows: + Since `pixel_values` is now a list with one tensor per image, we can override + [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows: ```python def _get_mm_fields_config( @@ -620,9 +493,15 @@ return a schema of the tensors outputted by the HF processor that are related to hf_inputs: BatchFeature, hf_processor_mm_kwargs: Mapping[str, object], ) -> Mapping[str, MultiModalFieldConfig]: - return dict(image_patches=MultiModalFieldConfig.batched("image")) + return dict( + pixel_values=MultiModalFieldConfig.batched("image"), + image_embeds=MultiModalFieldConfig.batched("image"), + ) ``` + !!! note + See our [actual code](../../../vllm/model_executor/models/mistral3.py) for the full implementation. + ### Prompt updates Override [_get_prompt_updates][vllm.multimodal.processing.BaseMultiModalProcessor._get_prompt_updates] to @@ -678,121 +557,53 @@ Each [PromptUpdate][vllm.multimodal.processing.PromptUpdate] instance specifies ] ``` -=== "Handling additional tokens: Fuyu" +=== "Handling additional tokens: PaliGemma" - Recall the layout of feature tokens from Step 2: - - ``` - |SPEAKER||SPEAKER|...|SPEAKER||NEWLINE| - |SPEAKER||SPEAKER|...|SPEAKER||NEWLINE| - ... - |SPEAKER||SPEAKER|...|SPEAKER||NEWLINE| - ``` - - We define a helper function to return `ncols` and `nrows` directly: + PaliGemma's HF processor inserts, after the prompt's leading `` token, a + run of image tokens followed by a second `` token that marks the start of + the text prompt. We start by building the run of image tokens, one per + placeholder feature token: ??? code ```python - def get_image_feature_grid_size( - self, - *, - image_width: int, - image_height: int, - ) -> tuple[int, int]: - image_processor = self.get_image_processor() - target_width = image_processor.size["width"] - target_height = image_processor.size["height"] - patch_width = image_processor.patch_size["width"] - patch_height = image_processor.patch_size["height"] - - if not (image_width <= target_width and image_height <= target_height): - height_scale_factor = target_height / image_height - width_scale_factor = target_width / image_width - optimal_scale_factor = min(height_scale_factor, width_scale_factor) - - image_height = int(image_height * optimal_scale_factor) - image_width = int(image_width * optimal_scale_factor) - - ncols = math.ceil(image_width / patch_width) - nrows = math.ceil(image_height / patch_height) - return ncols, nrows - ``` - - Based on this, we can initially define our replacement tokens as: - - ??? code - - ```python - def get_replacement(item_idx: int): - images = mm_items.get_items("image", ImageProcessorItems) - image_size = images.get_image_size(item_idx) - - ncols, nrows = self.info.get_image_feature_grid_size( - image_width=image_size.width, - image_height=image_size.height, + def get_insertion(item_idx: int): + images = mm_items.get_items( + "image", (ImageEmbeddingItems, ImageProcessorItems) ) - # `_IMAGE_TOKEN_ID` corresponds to `|SPEAKER|` - # `_NEWLINE_TOKEN_ID` corresponds to `|NEWLINE|` - return ([_IMAGE_TOKEN_ID] * ncols + [_NEWLINE_TOKEN_ID]) * nrows + if isinstance(images, ImageEmbeddingItems): + num_image_tokens = images.get_feature_size(item_idx) + else: + image_size = images.get_image_size(item_idx) + num_image_tokens = self.info.get_num_image_tokens( + image_width=image_size.width, + image_height=image_size.height, + ) + + image_tokens = [image_token_id] * num_image_tokens + ... ``` - However, this is not entirely correct. After `FuyuImageProcessor.preprocess_with_tokenizer_info` is called, - a BOS token (``) is also added to the prompt: + The trailing `` token is an additional token that must **not** receive a + vision embedding. To assign the vision embeddings to only the image tokens, + instead of returning the token ids directly you can return an instance of + [PromptUpdateDetails][vllm.multimodal.processing.PromptUpdateDetails] and mark + the embedding tokens with `embed_token_id`: ??? code ```python - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/fuyu/processing_fuyu.py#L417-L435 - model_image_input = self.image_processor.preprocess_with_tokenizer_info( - image_input=tensor_batch_images, - image_present=image_present, - image_unpadded_h=image_unpadded_heights, - image_unpadded_w=image_unpadded_widths, - image_placeholder_id=image_placeholder_id, - image_newline_id=image_newline_id, - variable_sized=True, - ) - prompt_tokens, prompts_length = _tokenize_prompts_with_image_and_batch( - tokenizer=self.tokenizer, - prompts=prompts, - scale_factors=scale_factors, - max_tokens_to_generate=self.max_tokens_to_generate, - max_position_embeddings=self.max_position_embeddings, - add_BOS=True, - add_beginning_of_answer_token=True, + return PromptUpdateDetails.select_token_id( + image_tokens + [bos_token_id], + embed_token_id=image_token_id, ) ``` - To assign the vision embeddings to only the image tokens, instead of a string - you can return an instance of [PromptUpdateDetails][vllm.multimodal.processing.PromptUpdateDetails]: - - ??? code - - ```python - hf_config = self.info.get_hf_config() - bos_token_id = hf_config.bos_token_id # `` - assert isinstance(bos_token_id, int) - - def get_replacement_fuyu(item_idx: int): - images = mm_items.get_items("image", ImageProcessorItems) - image_size = images.get_image_size(item_idx) - - ncols, nrows = self.info.get_image_feature_grid_size( - image_width=image_size.width, - image_height=image_size.height, - ) - image_tokens = ([_IMAGE_TOKEN_ID] * ncols + [_NEWLINE_TOKEN_ID]) * nrows - - return PromptUpdateDetails.select_token_id( - image_tokens + [bos_token_id], - embed_token_id=_IMAGE_TOKEN_ID, - ) - ``` - - Finally, noticing that the HF processor removes the `|ENDOFTEXT|` token from the tokenized prompt, - we can search for it to conduct the replacement at the start of the string: + Putting it together, we override [_get_prompt_updates][vllm.multimodal.processing.BaseMultiModalProcessor._get_prompt_updates]. + Since these tokens are inserted (rather than replacing an existing placeholder) + after the prompt's leading ``, we use [PromptInsertion][vllm.multimodal.processing.PromptInsertion] + with a prefix target: ??? code @@ -804,33 +615,41 @@ Each [PromptUpdate][vllm.multimodal.processing.PromptUpdate] instance specifies out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: hf_config = self.info.get_hf_config() - bos_token_id = hf_config.bos_token_id - assert isinstance(bos_token_id, int) + image_token_id = hf_config.image_token_index tokenizer = self.info.get_tokenizer() - eot_token_id = tokenizer.bos_token_id - assert isinstance(eot_token_id, int) - def get_replacement_fuyu(item_idx: int): - images = mm_items.get_items("image", ImageProcessorItems) - image_size = images.get_image_size(item_idx) + bos_token_id = tokenizer.bos_token_id + assert isinstance(bos_token_id, int) - ncols, nrows = self.info.get_image_feature_grid_size( - image_width=image_size.width, - image_height=image_size.height, + def get_insertion(item_idx: int): + images = mm_items.get_items( + "image", (ImageEmbeddingItems, ImageProcessorItems) ) - image_tokens = ([_IMAGE_TOKEN_ID] * ncols + [_NEWLINE_TOKEN_ID]) * nrows + + if isinstance(images, ImageEmbeddingItems): + num_image_tokens = images.get_feature_size(item_idx) + else: + image_size = images.get_image_size(item_idx) + num_image_tokens = self.info.get_num_image_tokens( + image_width=image_size.width, + image_height=image_size.height, + ) + + image_tokens = [image_token_id] * num_image_tokens return PromptUpdateDetails.select_token_id( image_tokens + [bos_token_id], - embed_token_id=_IMAGE_TOKEN_ID, + embed_token_id=image_token_id, ) return [ - PromptReplacement( + PromptInsertion( modality="image", - target=[eot_token_id], - replacement=get_replacement_fuyu, + target=PromptIndexTargets.prefix( + [bos_token_id] if tokenizer.add_bos_token else [] + ), + insertion=get_insertion, ) ] ``` @@ -873,7 +692,7 @@ Examples: Examples: - Chameleon (appends `sep_token`): [vllm/model_executor/models/chameleon.py](../../../vllm/model_executor/models/chameleon.py) -- Fuyu (appends `boa_token`): [vllm/model_executor/models/fuyu.py](../../../vllm/model_executor/models/fuyu.py) +- Molmo2 (prepends `bos_token`): [vllm/model_executor/models/molmo2.py](../../../vllm/model_executor/models/molmo2.py) - Molmo (applies chat template which is not defined elsewhere): [vllm/model_executor/models/molmo.py](../../../vllm/model_executor/models/molmo.py) ### Custom HF processor @@ -884,4 +703,3 @@ Examples: - DeepSeek-VL2: [vllm/model_executor/models/deepseek_vl2.py](../../../vllm/model_executor/models/deepseek_vl2.py) - InternVL: [vllm/model_executor/models/internvl.py](../../../vllm/model_executor/models/internvl.py) -- Qwen-VL: [vllm/model_executor/models/qwen_vl.py](../../../vllm/model_executor/models/qwen_vl.py) diff --git a/docs/deployment/frameworks/lws.md b/docs/deployment/frameworks/lws.md index 47586bcd700..5aae73c8a38 100644 --- a/docs/deployment/frameworks/lws.md +++ b/docs/deployment/frameworks/lws.md @@ -7,108 +7,202 @@ vLLM can be deployed with [LWS](https://github.com/kubernetes-sigs/lws) on Kuber ## Prerequisites -* At least two Kubernetes nodes, each with 8 GPUs, are required. -* Install LWS by following the instructions found [here](https://lws.sigs.k8s.io/docs/installation/). +- At least two Kubernetes nodes, each with 8 GPUs, are required. +- Install LWS by following the instructions found [here](https://lws.sigs.k8s.io/docs/installation/). ## Deploy and Serve -Deploy the following yaml file `lws.yaml` +Deploy the following yaml file `lws.yaml` (we have examples that use multiprocessing or Ray): -??? code "Yaml" +??? code "lws.yaml" + === "Multiprocessing (default)" + ```yaml + apiVersion: leaderworkerset.x-k8s.io/v1 + kind: LeaderWorkerSet + metadata: + name: vllm + spec: + replicas: 1 + leaderWorkerTemplate: + size: 2 + restartPolicy: RecreateGroupOnPodRestart + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - name: vllm-leader + image: docker.io/vllm/vllm-openai:latest + env: + - name: HF_TOKEN + value: + command: + - sh + - -c + - "vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size $(LWS_GROUP_SIZE) --nnodes $(LWS_GROUP_SIZE) --node-rank $(LWS_WORKER_INDEX) --master-addr $(LWS_LEADER_ADDRESS) --port 8080" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + ports: + - containerPort: 8080 + readinessProbe: + tcpSocket: + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + workerTemplate: + spec: + containers: + - name: vllm-worker + image: docker.io/vllm/vllm-openai:latest + command: + - sh + - -c + - "vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size $(LWS_GROUP_SIZE) --nnodes $(LWS_GROUP_SIZE) --node-rank $(LWS_WORKER_INDEX) --master-addr $(LWS_LEADER_ADDRESS) --headless" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + env: + - name: HF_TOKEN + value: + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + --- + apiVersion: v1 + kind: Service + metadata: + name: vllm-leader + spec: + ports: + - name: http + port: 8080 + protocol: TCP + targetPort: 8080 + selector: + leaderworkerset.sigs.k8s.io/name: vllm + role: leader + type: ClusterIP + ``` - ```yaml - apiVersion: leaderworkerset.x-k8s.io/v1 - kind: LeaderWorkerSet - metadata: - name: vllm - spec: - replicas: 1 - leaderWorkerTemplate: - size: 2 - restartPolicy: RecreateGroupOnPodRestart - leaderTemplate: - metadata: - labels: - role: leader - spec: - containers: - - name: vllm-leader - image: docker.io/vllm/vllm-openai:latest - env: - - name: HF_TOKEN - value: - command: - - sh - - -c - - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); - vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline_parallel_size 2" - resources: - limits: - nvidia.com/gpu: "8" - memory: 1124Gi - ephemeral-storage: 800Gi - requests: - ephemeral-storage: 800Gi - cpu: 125 - ports: - - containerPort: 8080 - readinessProbe: - tcpSocket: - port: 8080 - initialDelaySeconds: 15 - periodSeconds: 10 - volumeMounts: - - mountPath: /dev/shm - name: dshm - volumes: - - name: dshm - emptyDir: - medium: Memory - sizeLimit: 15Gi - workerTemplate: - spec: - containers: - - name: vllm-worker - image: docker.io/vllm/vllm-openai:latest - command: - - sh - - -c - - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(LWS_LEADER_ADDRESS)" - resources: - limits: - nvidia.com/gpu: "8" - memory: 1124Gi - ephemeral-storage: 800Gi - requests: - ephemeral-storage: 800Gi - cpu: 125 - env: - - name: HF_TOKEN - value: - volumeMounts: - - mountPath: /dev/shm - name: dshm - volumes: - - name: dshm - emptyDir: - medium: Memory - sizeLimit: 15Gi - --- - apiVersion: v1 - kind: Service - metadata: - name: vllm-leader - spec: - ports: - - name: http - port: 8080 - protocol: TCP - targetPort: 8080 - selector: - leaderworkerset.sigs.k8s.io/name: vllm - role: leader - type: ClusterIP - ``` + === "Ray" + ```yaml + apiVersion: leaderworkerset.x-k8s.io/v1 + kind: LeaderWorkerSet + metadata: + name: vllm + spec: + replicas: 1 + leaderWorkerTemplate: + size: 2 + restartPolicy: RecreateGroupOnPodRestart + leaderTemplate: + metadata: + labels: + role: leader + spec: + containers: + - name: vllm-leader + image: docker.io/vllm/vllm-openai:latest + env: + - name: HF_TOKEN + value: + command: + - sh + - -c + - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); + vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + ports: + - containerPort: 8080 + readinessProbe: + tcpSocket: + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + workerTemplate: + spec: + containers: + - name: vllm-worker + image: docker.io/vllm/vllm-openai:latest + command: + - sh + - -c + - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(LWS_LEADER_ADDRESS)" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + env: + - name: HF_TOKEN + value: + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + --- + apiVersion: v1 + kind: Service + metadata: + name: vllm-leader + spec: + ports: + - name: http + port: 8080 + protocol: TCP + targetPort: 8080 + selector: + leaderworkerset.sigs.k8s.io/name: vllm + role: leader + type: ClusterIP + ``` ```bash kubectl apply -f lws.yaml @@ -130,16 +224,37 @@ vllm-0-1 1/1 Running 0 2s Verify that the distributed tensor-parallel inference works: -```bash -kubectl logs vllm-0 |grep -i "Loading model weights took" -``` +=== "Multiprocessing (default)" + ```bash + kubectl logs vllm-0 | grep -i "Model loading" + kubectl logs vllm-0-1 | grep -i "Model loading" + ``` -Should get something similar to this: + Should get something similar to this: -```text -INFO 05-08 03:20:24 model_runner.py:173] Loading model weights took 0.1189 GB -(RayWorkerWrapper pid=169, ip=10.20.0.197) INFO 05-08 03:20:28 model_runner.py:173] Loading model weights took 0.1189 GB -``` + POD 0 (PP Rank 0) + + ```text + (Worker_PP0_TP0 pid=601) INFO 04-28 08:16:58 [gpu_model_runner.py:4820] Model loading took 3.82 GiB memory and 157.996399 seconds + ``` + + POD 1 (PP Rank 1) + + ```text + (Worker_PP1_TP0 pid=396) INFO 04-28 08:17:09 [gpu_model_runner.py:4820] Model loading took 3.82 GiB memory and 168.878781 seconds + ``` + +=== "Ray" + ```bash + kubectl logs vllm-0 | grep -i "Loading model weights took" + ``` + + Should get something similar to this: + + ```text + INFO 05-08 03:20:24 model_runner.py:173] Loading model weights took 0.1189 GB + (RayWorkerWrapper pid=169, ip=10.20.0.197) INFO 05-08 03:20:28 model_runner.py:173] Loading model weights took 0.1189 GB + ``` ## Access ClusterIP service @@ -173,7 +288,6 @@ curl http://localhost:8080/v1/completions \ The output should be similar to the following ??? console "Output" - ```text { "id": "cmpl-1bb34faba88b43f9862cfbfb2200949d", diff --git a/docs/deployment/integrations/kthena.md b/docs/deployment/integrations/kthena.md index 03ef190e558..7cc3f14a71e 100644 --- a/docs/deployment/integrations/kthena.md +++ b/docs/deployment/integrations/kthena.md @@ -64,36 +64,74 @@ A simplified version of the example (`llama-multinode`) looks like: - `spec.replicas: 1` – one `ServingGroup` (one logical model deployment). - `roles`: - `entryTemplate` – defines **leader** pods that run: - - vLLM’s **multi-node cluster bootstrap script** (Ray cluster). + - vLLM’s **multi-node cluster bootstrap script**. - vLLM **OpenAI-compatible API server**. - - `workerTemplate` – defines **worker** pods that join the leader’s Ray cluster. + - `workerTemplate` – defines **worker** pods to join the leader’s Ray cluster (Ray backend) or to join same distributed process group (multiprocessing backend). Key points from the example YAML: -- **Image**: `vllm/vllm-openai:latest` (matches upstream vLLM images). -- **Command** (leader): +Image: `vllm/vllm-openai:latest` (matches upstream vLLM images). +Commands: - ```yaml - command: - - sh - - -c - - > - bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2; - vllm serve meta-llama/Llama-3.1-405B-Instruct - --port 8080 - --tensor-parallel-size 8 - --pipeline-parallel-size 2 - ``` +??? code "Yaml" + === "Multiprocessing (default)" + Leader: -- **Command** (worker): + ```yaml + command: + - sh + - -c + - > + vllm serve meta-llama/Llama-3.1-405B-Instruct + --tensor-parallel-size 8 + --pipeline-parallel-size 2 + --nnodes=2 + --node-rank=0 + --master-addr=$(ENTRY_ADDRESS) + --port 8080 + ``` - ```yaml - command: - - sh - - -c - - > - bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS) - ``` + Worker: + + ```yaml + command: + - sh + - -c + - > + vllm serve meta-llama/Llama-3.1-405B-Instruct + --tensor-parallel-size 8 + --pipeline-parallel-size 2 + --nnodes=2 + --node-rank=1 + --master-addr=$(ENTRY_ADDRESS) + --headless + ``` + + === "Ray" + Leader: + + ```yaml + command: + - sh + - -c + - > + bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh + leader --ray_cluster_size=2; python3 -m + vllm.entrypoints.openai.api_server --port 8080 --model + meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 + --pipeline-parallel-size 2 + ``` + + Worker: + + ```yaml + command: + - sh + - -c + - > + bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh + worker --ray_address=$(ENTRY_ADDRESS) + ``` --- @@ -111,96 +149,192 @@ kubectl create secret generic hf-token \ ### 3.2 Apply the `ModelServing` +Save one of the following manifests to `modelserving.yaml`: + +??? code "modelserving.yaml" + === "Multiprocessing (default)" + ```yaml + apiVersion: workload.serving.volcano.sh/v1alpha1 + kind: ModelServing + metadata: + name: llama-multinode + namespace: default + spec: + schedulerName: volcano + replicas: 1 # group replicas + template: + restartGracePeriodSeconds: 60 + gangPolicy: + minRoleReplicas: + 405b: 1 + roles: + - name: 405b + replicas: 2 + entryTemplate: + spec: + containers: + - name: leader + image: vllm/vllm-openai:latest + env: + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HUGGING_FACE_HUB_TOKEN + command: + - sh + - -c + - "vllm serve meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2 --nnodes 2 --node-rank 0 --master-addr $(ENTRY_ADDRESS) --distributed-executor-backend mp --port 8080" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + ports: + - containerPort: 8080 + readinessProbe: + tcpSocket: + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + workerReplicas: 1 + workerTemplate: + spec: + containers: + - name: worker + image: vllm/vllm-openai:latest + command: + - sh + - -c + - "vllm serve meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2 --nnodes 2 --node-rank 1 --master-addr $(ENTRY_ADDRESS) --distributed-executor-backend mp --headless" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + env: + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HUGGING_FACE_HUB_TOKEN + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + ``` + + === "Ray" + ```yaml + apiVersion: workload.serving.volcano.sh/v1alpha1 + kind: ModelServing + metadata: + name: llama-multinode + namespace: default + spec: + schedulerName: volcano + replicas: 1 # group replicas + template: + restartGracePeriodSeconds: 60 + gangPolicy: + minRoleReplicas: + 405b: 1 + roles: + - name: 405b + replicas: 2 + entryTemplate: + spec: + containers: + - name: leader + image: vllm/vllm-openai:latest + env: + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HUGGING_FACE_HUB_TOKEN + command: + - sh + - -c + - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2; + vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + ports: + - containerPort: 8080 + readinessProbe: + tcpSocket: + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + workerReplicas: 1 + workerTemplate: + spec: + containers: + - name: worker + image: vllm/vllm-openai:latest + command: + - sh + - -c + - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS)" + resources: + limits: + nvidia.com/gpu: "8" + memory: 1124Gi + ephemeral-storage: 800Gi + requests: + ephemeral-storage: 800Gi + cpu: 125 + env: + - name: HUGGING_FACE_HUB_TOKEN + valueFrom: + secretKeyRef: + name: hf-token + key: HUGGING_FACE_HUB_TOKEN + volumeMounts: + - mountPath: /dev/shm + name: dshm + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 15Gi + ``` + ```bash -cat < **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. @@ -169,24 +159,39 @@ Priority is **1 = highest** (tried first). | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 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, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 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, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x | +| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 | +| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | 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_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | +| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | | `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 | -| `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 | +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any | +| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 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`. +> **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`. > > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise. +## MiniMax M3 Sparse Attention Backends + +Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer") +layers. It is wired in directly by the model and is not part of the +automatic priority lists above. A lightning indexer scores KV blocks, the +top-k blocks (plus fixed init/local blocks) are selected, and attention +attends only to those blocks; index keys live in a separate side cache. + +| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | +| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | +| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | + ## MLA (Multi-head Latent Attention) Backends MLA uses separate backends for prefill and decode phases. @@ -201,9 +206,9 @@ hardware and configuration. | Backend | Description | Dtypes | Compute Cap. | Notes | | ------- | ----------- | ------ | ------------ | ----- | | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise | -| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only | -| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only | -| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only | +| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only | +| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | +| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | > **‡** Automatic selection tries FlashAttention first. On Blackwell > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then @@ -217,11 +222,13 @@ MLA decode backends are selected using the standard | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `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 | +| `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 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.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 | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | +| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x | | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | @@ -235,10 +242,11 @@ DeepSeek V4 sparse MLA uses its own decode backends, selected via `--attention-backend=` (e.g., `FLASHMLA_SPARSE_DSV4`, `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index pipeline (compressor + SWA + indexer, 256-token blocks, head 512); -default on NVIDIA is `FLASHMLA_SPARSE_DSV4`. +default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and +`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures. | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any | -| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | +| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x | +| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index 718a4a8154d..e274b68c702 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -161,11 +161,11 @@ class AttentionCGSupport(enum.Enum): ALWAYS = 3 """CUDA Graphs always supported; supports mixed-prefill-decode""" UNIFORM_BATCH = 2 - """CUDA Graphs supported for batches the only contain query lengths that are + """CUDA Graphs supported for batches that only contain query lengths that are the same, this can be used for spec-decode i.e. "decodes" are 1 + num_speculative_tokens""" UNIFORM_SINGLE_TOKEN_DECODE = 1 - """CUDA Graphs supported for batches the only contain query_len==1 decodes""" + """CUDA Graphs supported for batches that only contain query_len==1 decodes""" NEVER = 0 """NO CUDA Graphs support""" ``` diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 8cbbedf9d0b..7eab425d6e2 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -2,6 +2,8 @@ The [CUDA Graphs](cuda_graphs.md) infrastructure in vLLM primarily targets the **decoder** (language model) forward pass. vLLM also supports capturing the **encoder** (vision transformer) forward pass as CUDA Graphs, independently from the decoder. This is based on . +For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tiling), a **dual-path graph** mode captures two independent sets of CUDA graphs — one for the global image path and one for the local patch path — enabling independent budget selection and partial eager fallback per path. This is based on . + !!! note Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md). @@ -11,6 +13,8 @@ Vision encoder inference incurs CUDA kernel launch overhead on the host side. Th Encoder CUDA Graphs eliminate this overhead by pre-capturing the full encoder forward pass at multiple token budget levels during model initialization, then replaying the appropriate graph at runtime. +For two-tower vision encoders such as DeepSeek-OCR (SAM + CLIP with dynamic tiling), the global image path and local patch path have independent token profiles (272 tokens per global image vs. 100 tokens per local patch). Capturing a single monolithic graph for both paths would significantly reduce packing efficiency. The dual-path graph mode captures each path as a separate set of budgets, allowing the manager to pack and replay each path independently. + ## Design The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, managed by [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]. The system contains the following core components: @@ -37,10 +41,14 @@ class BudgetGraphMetadata: Budgets are auto-generated as power-of-2 levels from a model-provided range via `get_encoder_cudagraph_budget_range()`, with the maximum budget always included even if it does not fall on a power-of-2 boundary. Budgets can also be explicitly specified by the user via `encoder_cudagraph_token_budgets` in `CompilationConfig`. +When `EncoderCudaGraphConfig.enable_dual_path_graph` is `True`, the manager generates two independent budget lists — `global_token_budgets` (multiples of `global_token_per_image`) and `local_token_budgets` (multiples of `local_token_per_patch`) — and stores captured graphs under `budget_graphs["global"]` and `budget_graphs["local"]` respectively. + ### Greedy bin-packing at runtime When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. +For dual-path models, the manager routes to `_execute_local_dual_path()`, which constrains both global and local token budgets simultaneously during packing (see [Dual-Path graph capture](#dual-path-graph-capture)). + For each graph replay: 1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. @@ -48,6 +56,42 @@ For each graph replay: 3. Replay the CUDA Graph. 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). +### Dual-Path graph capture + +For two-tower vision encoders (e.g., DeepSeek-OCR), the `EncoderCudaGraphConfig` sets `enable_dual_path_graph=True` and provides `global_token_per_image` / `local_token_per_patch`. The manager captures two independent sets of CUDA graphs — one for the **global** image path and one for the **local** patch path — stored under `budget_graphs["global"]` and `budget_graphs["local"]` respectively. + +**Budget generation.** Two separate budget lists are generated: + +* `global_token_budgets` — power-of-2 multiples of `global_token_per_image` (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR). +* `local_token_budgets` — power-of-2 multiples of `local_token_per_patch` (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800]` for DeepSeek-OCR). A budget of `0` is always included to handle images with no local patches (images ≤ 640×640 that produce only global features). + +Both lists are capped at the same `max_budget`. + +**Dual-path greedy packing.** Each `EncoderItemSpec` provides both `global_output_tokens` (constant per image) and `local_output_tokens` (proportional to the patch count). The dual-path packing algorithm constrains both budgets simultaneously: + +* Sort images by total output tokens (global + local), smallest first. +* Greedily pack images: an image is added to the current sub-batch only if both the accumulated global tokens ≤ `max_global_budget` **and** the accumulated local tokens ≤ `max_local_budget`, with the image count ≤ `max_batch_size`. +* Once either constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path. +* Repeat until all images are packed. + +**Partial graph fallback.** After packing, each sub-batch falls into one of four execution scenarios: + +| Global budget | Local budget | Execution | +| :---: | :---: | --- | +| Found | Found | Both paths use CUDA graph replay | +| Found | `None` | Global graph replay + local path skipped (no patches) | +| `None` | Found | Global eager fallback + local graph replay | +| `None` | `None` | Both paths fall back to eager execution | + +Note that the `0`-budget graph is never actually replayed for local — it signals that local patch processing should be skipped entirely. + +**Buffer keys per path.** Global and local paths use different buffer keys. For DeepSeek-OCR, the global path uses `pixel_values` (full images, shape `[B, 3, 1280, 1280]`) while the local path uses `images_crop` (patches, shape `[P, 3, 1024, 1024]`). The manager iterates over each captured graph's own `input_buffers.keys()` rather than a shared `buffer_keys` list, so both paths can use different buffers. + +**Post-processing.** The `postprocess_encoder_output` method receives a `local_output` parameter (a tensor or `None`) containing the local-path encoder output. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, this means reshaping the global output into `[B, 272, n_embed]`, the local output into `[P, 100, n_embed]`, assembling patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image. + +!!! note + The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image. + ### Data-parallel support 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`. @@ -67,28 +111,33 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra * `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video). * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets. -* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size and output token count. Replaces the former three separate methods (`get_num_items`, `get_per_item_output_tokens`, `get_per_item_input_sizes`). +* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and optionally per-path token counts (`global_output_tokens`, `local_output_tokens`) for dual-path models. * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding. -* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph. -* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch)` — computes buffer values from actual batch inputs. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match `buffer_keys` in the config. -* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor])` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `pixel_values` tensor is included in `inputs` alongside metadata buffers. -* `encoder_eager_forward(mm_kwargs)` — fallback eager forward when no graph fits. -* `postprocess_encoder_output(...)` — post-process encoder output, delegates to `scatter_output_slices` by default. +* `prepare_encoder_cudagraph_capture_inputs(..., path="default")` — creates dummy inputs for graph capture. The `path` parameter (`"global"` or `"local"`) tells the model which path to generate dummy inputs for. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph. +* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch, path="default")` — computes buffer values from actual batch inputs. The `path` parameter selects which modality keys to extract from `mm_kwargs`. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match the captured graph's `input_buffers.keys()`. +* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor], path="default")` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `path` parameter dispatches to the correct encoder sub-module (e.g., global vs. local path for DeepSeek-OCR). +* `encoder_eager_forward(mm_kwargs, path="default")` — fallback eager forward when no graph fits. When `path` is `"global"` or `"local"`, runs only that encoder path without graph capture. +* `postprocess_encoder_output(..., local_output=None)` — post-process encoder output. The `local_output` parameter receives the local-path encoder output tensor (or `None`), enabling dual-path models to assemble global and local features into the final per-image embedding. !!! 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 | -| ------------ | ------ | ------------ | ------------ | -| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | -| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | -| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | -| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | -| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | -| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | -| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | +| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph | +| ------------ | ------ | ------------ | ------------ | --------------- | +| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | +| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | +| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | +| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | +| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | +| `Step3VLForConditionalGeneration` | `Step3-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. @@ -103,6 +152,8 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs: * `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 from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). 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). +Dual-path mode is configured at the model level via `EncoderCudaGraphConfig` fields (`enable_dual_path_graph`, `global_token_per_image`, `local_token_per_patch`) — no additional user configuration is required. The manager automatically generates separate budget lists and routes to dual-path execution when the model opts in. + ## Usage guide ### Image inference @@ -114,6 +165,14 @@ vllm serve Qwen/Qwen3-VL-32B \ --compilation-config '{"cudagraph_mm_encoder": true}' ``` +For `Llama 4` (image only): + +```bash +vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \ + --limit-mm-per-prompt '{"image": 1}' \ + --compilation-config '{"cudagraph_mm_encoder": true}' +``` + With explicit budgets: ```bash diff --git a/docs/design/endpoint_plugins.md b/docs/design/endpoint_plugins.md new file mode 100644 index 00000000000..9f38fe5da18 --- /dev/null +++ b/docs/design/endpoint_plugins.md @@ -0,0 +1,136 @@ +# Endpoint Plugins + +Endpoint plugins let out-of-tree packages add HTTP routes to the OpenAI compatible API server without editing `vllm/entrypoints/openai/api_server.py`. Their scope is +the **HTTP surface only** registering routes and optionally per app state used by those routes. A plugin reaches the engine the same way an in-tree serving handler does, through the `EngineClient` it is handed at startup (e.g. `engine_client.collective_rpc(...)`). No new engine access path is introduced. + +!!! warning "Security" + Endpoint plugins are **not loaded by default** and must be explicitly allowlisted. Read [Endpoint Plugins security posture](../usage/security.md#endpoint-plugins) before enabling one, especially the route shadowing warning. + +## The `EndpointPlugin` protocol + +Endpoint plugins implement the [`EndpointPlugin`][vllm.plugins.endpoint_plugins.interface.EndpointPlugin] runtime checkable `Protocol`: + +```python +class EndpointPlugin(Protocol): + name: str + required_tasks: tuple[SupportedTask, ...] | None + + def attach_router(self, app: FastAPI) -> None: ... + + async def init_state( + self, engine_client: EngineClient | None, state: State, args: Namespace + ) -> None: ... +``` + +- `name`: a unique identifier used in logs and for `VLLM_PLUGINS` allowlisting +- `required_tasks`: the tasks the server must support for this plugin to load. `None` means the plugin has no task requirement +- `attach_router`: registers routes on `app` +- `init_state`: initializes per app state the routes read at request time + +## The two phase lifecycle + +Routes are registered before the engine exists. This means the interface has to expose two hooks that run at two different points in server startup: + +| Phase | Called from | `engine_client` available? | Work | +| --- | --- | --- | --- | +| A. Route registration | `build_app()` | No | `attach_router(app)` add routes. Do not touch the engine here. | +| B. State init | `init_app_state()` | Usually but `None` on the CPU only render server | `init_state(engine_client, state, args)` build a serving handler holding `engine_client` and store it on `state`. | + +Because `app.state` *is* the `state` object passed to `init_app_state()`, an object stored during phase A is visible in phase B and an object stored in phase B is visible to route handlers at request time via `request.app.state`. This is the same pattern in-tree endpoints already use. + +### Engine less servers (the render server) + +The CPU only render server (`init_render_app_state()`) has no `EngineClient`. It still runs both phases for any plugin eligible for the `render` task (`required_tasks` is `None` or includes `"render"`). `attach_router` is called as usual but `init_state` is called with `engine_client=None`. + +A plugin that needs an engine to function has two options: + +- Exclude `"render"` from `required_tasks` so it is never loaded on the render server in the first place +- Accept being loaded on `render` and check for `None` in `init_state` or in the route handler returning an error response (e.g. HTTP 503) instead of dereferencing a client that doesn't exist + +`tests/plugins/vllm_add_dummy_endpoint_plugin` demonstrates the second option. Its route handler returns a 503 when `state.dummy_engine_client` is `None`. + +### Reaching the engine from a route handler + +`init_state` is where a plugin captures `engine_client` into a small serving handler and stashes it on `state`. The route added in `attach_router` reads that handler off `request.app.state` at request time and calls the engine through it, typically via `engine_client.collective_rpc(...)`. + +This minimal example omits the `None` check from the previous section for brevity since `required_tasks` is `None` here. It is in fact eligible for `render` and should handle `engine_client=None` the way `tests/plugins/vllm_add_dummy_endpoint_plugin` does before shipping it: + +```python +from fastapi import FastAPI, Request + + +class MyAdminEndpointPlugin: + name = "my_admin_endpoint_plugin" + required_tasks: tuple[str, ...] | None = None + + def attach_router(self, app: FastAPI) -> None: + @app.get("/plugins/my_admin_endpoint_plugin/scheduler_config") + async def scheduler_config(raw_request: Request): + engine_client = raw_request.app.state.my_engine_client + results = await engine_client.collective_rpc("get_scheduler_config") + return {"scheduler_config": results} + + async def init_state(self, engine_client, state, args) -> None: + state.my_engine_client = engine_client +``` + +A complete and tested version of this example is in-repo as `tests/plugins/vllm_add_dummy_endpoint_plugin` and is exercised e2e (including a real HTTP request) in `tests/plugins_tests/test_endpoint_plugins.py`. + +## Registering the entry point + +Register a zero argument factory (a class or function) under the `vllm.endpoint_plugins` group. The factory must return an object satisfying `EndpointPlugin`: + +```toml +# pyproject.toml +[project.entry-points."vllm.endpoint_plugins"] +my_admin_api = "my_pkg.endpoints:MyAdminEndpointPlugin" +``` + +```python +# setup.py equivalent +setup( + name="my_pkg", + entry_points={ + "vllm.endpoint_plugins": [ + "my_admin_api = my_pkg.endpoints:MyAdminEndpointPlugin" + ] + }, +) +``` + +The entry point name (`my_admin_api` above) is independent of the plugin's `name` attribute. `VLLM_PLUGINS` allowlisting matches on the **entry point name** following the same convention as `vllm.general_plugins` (see [Plugin System](plugin_system.md)). + +## Gating: `VLLM_PLUGINS` and `required_tasks` + +Endpoint plugins are discovered and gated by [`load_endpoint_plugins`][vllm.plugins.load_endpoint_plugins] which is stricter than the loader used for other plugin groups: + +- **Nothing loads unless `VLLM_PLUGINS` is set and names the plugin.** Other plugin groups load everything unless `VLLM_PLUGINS` narrows the set. Endpoint plugins invert that default because they add network exposed surface. See [Security](../usage/security.md#endpoint-plugins). +- **`required_tasks` must intersect the server's supported tasks** unless it is `None`. Use this to keep a plugin from attaching routes on a server that can't service them (e.g. a pooling only deployment). +- A factory that raises an issue during instantiation is logged and skipped. It does not abort server startup. + +Only the front end API server process loads endpoint plugins. There is no need to guard for worker or engine core processes. + +## Pairing with `vllm.general_plugins` + +Endpoint plugins cover the HTTP surface only. If a plugin also needs new engine side behavior (a new worker-side RPC method, a custom stat) that half ships separately through the existing `vllm.general_plugins` group which loads in worker processes (see [Plugin System](plugin_system.md)). The two entry points are registered and loaded **independently**. Neither implies the other. The recommended distribution shape is a single package exposing both: + +```toml +[project.entry-points."vllm.general_plugins"] +my_admin_engine = "my_pkg.engine:register" # adds the worker side method + +[project.entry-points."vllm.endpoint_plugins"] +my_admin_api = "my_pkg.endpoints:MyAdminEndpointPlugin" # adds the HTTP route +``` + +Do not expect a single endpoint plugin to also mutate engine/worker state. If your route needs a worker side method that doesn't already exist then add it via a paired `general_plugins` entry point. + +## Path-prefix convention + +There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). A plugin's `attach_router` can register a path that collides with a core route and routes attached later win. To avoid surprising operators: + +- Namespace your routes under a distinct prefix, e.g. `/plugins//...`, rather than reusing `/v1/...` or other core prefixes +- Only register routes under a core prefix (like the worked example's `/v1/admin/scheduler_config`) if you specifically intend to override or extend existing behavior and document that clearly for operators allowlisting your plugin + +## Compatibility + +`state`/serving handler internals (e.g. the shape of in-tree `OpenAIServing*` classes) are not a stable public contract yet. Treat them as use-at-your-own-risk and expect them to change between vLLM versions. `FastAPI`, `EngineClient` and the `EndpointPlugin` protocol itself are the supported surface. diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 371a9c59320..c9991f75cdb 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -30,7 +30,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 | +| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm (+ FP8 quant) → 1 kernel | O1 (ROCm/AITER only) | 1-2% | No | Always | ## Support Matrix @@ -381,11 +381,32 @@ q_normed, kv_normed = fused_mla_dual_rms_norm( Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above when AITER is available. +**FP8 attention variant (per-token quant).** With a per-token FP8 `q_b_proj`, +only the *q* latent is FP8-quantized while *kv* stays bf16. +`RocmAiterRMSNormQuantFusionPass` first folds the q side into +`rocm_aiter_rmsnorm_fused_dynamic_quant`, leaving kv a plain +`rms_norm` — breaking the symmetric pattern above. The same pass then matches +this asymmetric pair and lowers it to `fused_mla_dual_rms_norm_per_token_quant`. + +```text +# Unfused (q norm+quant fused; kv still plain rms_norm): +q_c, kv_lora = split(projected, [q_dim, kv_dim]) +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim]) +q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8) +kv_normed = rms_norm(kv_c, kv_weight, eps) # bf16 + +# 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_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant( + q_c, q_weight, kv_c, kv_weight, eps1, eps2) +``` + **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) +- 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`, `MLADualRMSPerTokenQuantPattern`) +- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`, `fused_mla_dual_rms_norm_per_token_quant`) +- AITER kernels: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442), `fused_qk_rmsnorm_per_token_quant` ## See Also diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 8f17b473adc..82d54e9b5c1 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -159,7 +159,7 @@ For simplicity, we assume `block_size=1` in this section. ### High level idea -The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to catch the full blocks. That means the same tokens of different groups are cached and evicted independently. +The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to cache the full blocks. That means the same tokens of different groups are cached and evicted independently. When a new request comes in, we check the cache hit prefix of each group, and return the intersection of these groups as the cached prefix of the request. See below for the detailed algorithm for checking the cache hit of one group & performing the intersection. diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 0ae42039976..7b463b8750c 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -685,7 +685,7 @@ documentation for this option states: > use of possibly costly and or blocking operations and hence might > have a performance impact. -The metrics were added by and who up in an OpenTelemetry trace +The metrics were added by and show up in an OpenTelemetry trace as: ```text diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index af7da63b550..07d2a539801 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -42,7 +42,7 @@ th { 1. All types: mxfp4, nvfp4, int4, int8, fp8 2. A,T quantization occurs after dispatch. 3. All quantization happens after dispatch. - 4. Controlled by different env vars (`VLLM_FLASHINFER_MOE_BACKEND` "throughput" or "latency") + 4. Controlled by `--moe-backend` (`flashinfer_cutlass` or `flashinfer_trtllm`) 5. This is a no-op dispatcher that can be used to pair with any modular experts to produce a modular kernel that runs without dispatch or combine. These cannot be selected via environment variable. These are generally use for testing or adapting an expert subclass to the `fused_experts` API. 6. This depends on the experts implementation. @@ -89,6 +89,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | 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.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.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] | +| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md new file mode 100644 index 00000000000..b99ba6659f7 --- /dev/null +++ b/docs/design/nixl_kv_push_connector.md @@ -0,0 +1,256 @@ +# NIXL push-mode KV transfer + +The default NIXL connector is **pull-based**: the decode (D) instance +reads KV blocks from the prefill (P) instance via `NIXL READ` after +prefill completes. `NixlPushConnector` adds a **push-based** alternative +in which P writes the KV blocks directly into D's pre-allocated memory +via `NIXL WRITE`. + +This document describes the threading, queues, and scheduling +interactions specific to the push design. The pull-mode design is +unchanged; the push connector reuses the same handshake, NIXL agent +setup, and metadata path wherever possible. + +## High-level flow + +```mermaid +sequenceDiagram + autonumber + participant Client + participant Proxy + participant DSched as D Scheduler + participant DWorker as D Worker (main) + participant DWriter as D Writer + participant PWriter as P Writer + participant PWorker as P Worker (main) + participant PSched as P Scheduler + + Client->>Proxy: POST /v1/completions + Proxy->>PSched: prefill leg (do_remote_decode=True, max_tokens=1) + Proxy->>DSched: decode leg (do_remote_prefill=True, P coordinates) + + note over DSched,DWriter: D side - register blocks with P + DSched->>DSched: update_state_after_alloc, stash registration, arm watchdog + DSched->>DWorker: build_connector_meta -> meta.push_registrations + DWorker->>DWriter: enqueue (req_id, reg_data) on _reg_send_inbox + DWriter->>PWriter: NIXL send_notif PUSH_REG msgpack + + note over PSched,PWriter: P side - prefill, stage finished blocks + PSched->>PSched: request_finished, stash blocks + PSched->>PWorker: build_connector_meta -> meta.push_finished_blocks + PWorker->>PWriter: enqueue (req_id, blocks) on _finished_blocks_inbox + + note over PWriter: P writer matches and WRITEs + PWriter->>PWriter: get_new_notifs returns PUSH_REG, route via _handle_push_reg_notif + alt PUSH_REG and finished blocks both present + PWriter->>PWriter: pop matching pair, fire WRITE + else only one side present + PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched + end + PWriter->>PWriter: ensure D handshake (one-time) + PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif + + note over DWorker,DWriter: D side - completion accounting + DWriter-->>DWorker: forward HB and completion notifs via _pending_completion_notifs + DWorker->>DWorker: _get_new_notifs drains, HB extends lease, completion marks recv done + DWorker->>DSched: update_connector_output(finished_recving) + DSched->>DSched: clear watchdog deadline + + note over PWorker,PWriter: P side - reclaim + PWorker->>PWorker: get_finished, drain _sending_transfers, queue eviction + PWriter->>PWriter: drain _evict_finished_inbox, drop stale state + PWorker->>PSched: update_connector_output(finished_sending) + PSched->>PSched: free lease + + DWorker-->>Proxy: stream decode tokens + Proxy-->>Client: response +``` + +## Threads + +``NixlPushConnectorWorker`` introduces a single dedicated background +thread per worker (i.e. per TP rank), named ``nixl-push-writer``. +Each owns the new push-specific NIXL operations on its rank: + +* ``nixl_wrapper.get_new_notifs()`` — receive notifications. +* ``nixl_wrapper.send_notif(...)`` for the ``PUSH_REG:`` (D + side) and for the per-WRITE completion notif (P side). +* ``nixl_wrapper.make_prepped_xfer(...) / transfer(...)`` — submit the + WRITE itself. + +Heartbeats continue to go out from the engine main thread via the +existing base-worker ``_send_heartbeats`` plumbing inside +``start_load_kv``. + +### Wake model + +The writer thread blocks on ``_push_writer_wake`` (a +``threading.Event``) when it has no work. Three callers set the +event: + +1. **``start_load_kv``** (worker main thread, called once per engine + step with the scheduler's metadata) — sets the wake only when the + step actually hands the writer new work, i.e. when + ``meta.push_registrations`` or ``meta.push_finished_blocks`` is + non-empty. This is the wake for new transfers. +2. **``get_finished``** (worker main thread, called once per engine + step to report completions) — always sets the wake. The writer is + the sole consumer of ``nixl_wrapper.get_new_notifs()`` for push, + so this gives it a chance to drain inbound notifs (heartbeats from + D, completion notifs after a WRITE, late-arriving ``PUSH_REG``) + even when there is no new metadata to act on. +3. **Handshake-completion callback** (background handshake executor + thread) — when a deferred D→P handshake finishes successfully, the + future's done-callback re-enqueues the registration onto + ``_reg_send_inbox`` and sets the wake so the corresponding + ``send_notif`` runs on the writer (we never call ``send_notif`` from + the executor thread). On this second pass ``_ensure_handshake`` + returns ``None`` (the agent is now connected), so the writer sends + the ``PUSH_REG`` directly. If the handshake *failed*, the callback + fails the request instead of re-enqueuing, so there is no retry + loop. + +In addition to event-driven wakes, the writer self-polls at +``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side +finished blocks waiting for an unmatched ``PUSH_REG``. + +When a request completes on P (lease expires or the WRITE finishes), +``get_finished`` enqueues the request id onto ``_evict_finished_inbox``, +which the writer drains to drop stale ``_push_finished_blocks`` / +``_pending_d_registrations`` and stop self-polling. + +## Writer-local matching tables + +| Table | Owner | Holds | +|--------------------------------|------------------|------------------------------------------------------------------------| +| `_pending_d_registrations` | writer | D registrations received from a remote D, waiting for P's blocks | +| `_push_finished_blocks` | writer | P blocks staged by the scheduler, waiting for a remote D registration | + +Either side can arrive first. The writer matches in both directions: +when a ``PUSH_REG`` arrives we look up ``_push_finished_blocks``, and +when finished blocks arrive we look up ``_pending_d_registrations``. +Both lookups try an exact ``request_id`` match first, then fall back +to comparing the ids after stripping the trailing per-engine random +suffix (via ``get_base_request_id``). The fallback exists because the +proxy hands the same ``X-Request-Id`` to both legs, so P and D wrap it +into the same ``cmpl--`` form and differ only by the +8-hex randomization suffix that ``input_processor.assign_request_id`` +appends per engine. Stripping just that suffix normalizes both sides +to the same id while preserving the completion index (so multi-prompt +sub-requests stay distinct). It also works whether or not +``VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`` is set, which matters since +that env var is slated for removal upstream. + +## Wire format + +A push registration is sent as a NIXL notification: + +```text +PUSH_REG: +``` + +Fields in the dict: + +| Field | Set by | Meaning | +|----------------------|--------|------------------------------------------------------------------------| +| ``request_id`` | D | D's own vLLM request id; P's match key, echoed in the completion notif | +| ``decode_engine_id`` | D | D's engine id (P uses this for the reverse handshake) | +| ``decode_host`` | D | D's NIXL side-channel host | +| ``decode_port`` | D | D's NIXL side-channel port | +| ``decode_tp_size`` | D | D's tensor-parallel size | +| ``local_block_ids`` | D | per-group lists of D's *logical* block ids (preallocated) | +| ``remote_engine_id`` | D | P's engine id (for the existing P-side handshake) | +| ``remote_host`` | D | P's NIXL side-channel host | +| ``remote_port`` | D | P's NIXL side-channel port | +| ``remote_tp_size`` | D | P's tensor-parallel size | + +D ships **logical** block ids; P expands them to physical block ids at +WRITE-submission time using the ratio learned during the NIXL +handshake (`remote_physical_blocks_per_logical`). This matches the +pull-mode contract — schedulers ship logical ids, workers expand to +physical at submission. + +The completion notif sent from P to D after a WRITE is the existing +`:` format used in pull mode (here ``request_id`` +is D's own request id, taken from the registration), so the D-side +accounting code is unchanged. + +## Scheduler-side responsibilities + +`NixlPushConnectorScheduler` extends the base scheduler with: + +* **D side** — `update_state_after_alloc` stashes registration data in + `_push_pending_registrations` and arms a soft watchdog + (`_push_registration_deadlines`). `build_connector_meta` drains the + stash into `meta.push_registrations` and any expired entries are + dropped with a warning. +* **P side** — `request_finished` stashes block IDs in + `_finished_request_blocks` (for the lease and for + `has_pending_push_work`) and `_newly_finished_push_blocks` (for the + next worker step via `meta.push_finished_blocks`). +* **Both sides** — `has_pending_push_work` keeps the engine main loop + stepping while there is in-flight push state, so the writer always + gets at least one wake per step. + +`update_connector_output`: + +* `finished_sending` (P side) clears the lease entry. +* `finished_recving` (D side) clears the watchdog deadline. + +## Timeouts and watchdogs + +Two per-request timers are armed on the scheduler: + +* **D-side registration watchdog** — ``_push_registration_deadlines``. + If a registered request does not see a push completion within + ``push_registration_timeout`` seconds (defaults to + ``decoder_kv_blocks_ttl``), ``build_connector_meta`` drops the stale + registration and the pending entry, logs a warning, and stops trying + to resend the registration. The corresponding request remains tracked + in ``_reqs_need_recv``; it is the engine's request-level abort path + (or the user / proxy timing out the HTTP call) that ultimately fails + the request. +* **P-side block lease** — same ``_kv_lease_duration`` used by pull + mode. ``request_finished`` sets the expiration in ``_reqs_need_send`` + and ``update_connector_output(finished_sending=...)`` clears it on + successful WRITE. Stale leases are reaped by ``get_finished`` in the + base worker, which then enqueues the eviction onto + ``_evict_finished_inbox`` so the writer also stops self-polling. + +## Failure handling + +* **D-side handshake failure (P→D handshake before sending PUSH_REG)** — + the future's done-callback calls ``_handle_failed_transfer(rid, None)``, + which marks D's pre-allocated blocks invalid and enqueues onto + ``_failed_recv_reqs`` so the next ``get_finished`` reports the + request as a failed recv. Same recv-side accounting as pull mode. +* **D-side ``send_notif`` failure when shipping the PUSH_REG to P** — + identical handling: ``_handle_failed_transfer`` marks the recv as + failed. +* **P-side WRITE submission failure** — the WRITE handle (if any) is + released and ``xfer_stats.record_failed_transfer()`` bumps the + failure counter. We deliberately do not call + ``_handle_failed_transfer`` here: ``req_id`` on the P side has no + entry in ``_recving_metadata`` (P is not the receiver), so the + helper would put a P-local request id into ``_failed_recv_reqs`` + and trip the assertion in the base worker's ``get_finished``. The + outbound WRITE is dropped on the floor; D's lease watchdog handles + the missing completion. + +## Summary + +The push design is a small, well-contained extension on top of the +existing NIXL connector: + +* one new connector class, one new scheduler class, one new worker + class — all subclasses of the existing base classes; +* one dedicated background thread per worker; +* a few cross-thread queues, each with a single consumer (the writer); + most have one producer, except ``_reg_send_inbox``, which is fed both + by the engine main thread (new registrations) and by the + handshake-completion callback (registrations replayed after their + D→P handshake finishes); +* one new notification type (`PUSH_REG:`). + +Behavior on the engine main thread is otherwise unchanged. The writer +thread is event-driven and idle when there is no push work. diff --git a/docs/design/paged_attention.md b/docs/design/paged_attention.md index 7c0132cd2a2..f4742c7faaa 100644 --- a/docs/design/paged_attention.md +++ b/docs/design/paged_attention.md @@ -52,7 +52,7 @@ __device__ void paged_attention_kernel( ) ``` -There are also a list of template arguments above the function +There is also a list of template arguments above the function signature that are determined during compilation time. `scalar_t` represents the data type of the query, key, and value data elements, such as FP16. `HEAD_SIZE` indicates the number of elements in each @@ -178,7 +178,7 @@ const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride + physical_block_offset * x; ``` -Unlike to `q_ptr`, `k_ptr` in each thread will point to different +Unlike `q_ptr`, `k_ptr` in each thread will point to different key token at different iterations. As shown above, that `k_ptr` points to key token data based on `k_cache` at assigned block, assigned head and assigned token. diff --git a/docs/design/plugin_system.md b/docs/design/plugin_system.md index e5c9cea17c2..dd49df0ef2f 100644 --- a/docs/design/plugin_system.md +++ b/docs/design/plugin_system.md @@ -53,6 +53,8 @@ Every plugin has three parts: - **Stat logger plugins** (with group name `vllm.stat_logger_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree loggers into vLLM. The entry point should be a class that subclasses StatLoggerBase. +- **Endpoint plugins** (with group name `vllm.endpoint_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree HTTP routes on the OpenAI compatible API server. Unlike the other plugin groups above, endpoint plugins are loaded only in the API server front end process and are **not loaded by default**. See [Endpoint Plugins](endpoint_plugins.md) for the interface and [Security](../usage/security.md#endpoint-plugins) for the opt-in and trust model. + ## Guidelines for Writing Plugins - **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes. diff --git a/docs/design/prefix_caching.md b/docs/design/prefix_caching.md index 0f3100c9b73..f783f4a1bc8 100644 --- a/docs/design/prefix_caching.md +++ b/docs/design/prefix_caching.md @@ -27,7 +27,7 @@ In the example above, the KV cache in the first block can be uniquely identified For `vllm serve`, you can control the hashing algorithm via `--prefix-caching-hash-algo`: - `sha256` (default): Uses Python's `pickle` for serialization. Hashes may not be reproducible across different Python or vLLM versions. - `sha256_cbor`: Uses `cbor2` for serialization, providing a reproducible, cross-language compatible hash. This is recommended for deterministic caching across environments. - - `xxhash`: `Uses Pickle serialization with xxHash (128-bit) for faster, non-cryptographic hashing. Requires the optional `xxhash` package. IMPORTANT: Use of a hashing algorithm that is not considered cryptographically secure theoretically increases the risk of hash collisions, which can cause undefined behavior or even leak private information in multi-tenant environments. Even if collisions are still very unlikely, it is important to consider your security risk tolerance against the performance benefits before turning this on. + - `xxhash`: Uses Pickle serialization with xxHash (128-bit) for faster, non-cryptographic hashing. Requires the optional `xxhash` package. IMPORTANT: Use of a hashing algorithm that is not considered cryptographically secure theoretically increases the risk of hash collisions, which can cause undefined behavior or even leak private information in multi-tenant environments. Even if collisions are still very unlikely, it is important to consider your security risk tolerance against the performance benefits before turning this on. - `xxhash_cbor` combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional `xxhash` package. **A hashing example with multi-modality inputs** @@ -197,7 +197,7 @@ As can be seen, block 3 is a new full block and is cached. However, it is redund When a request is finished, we free all its blocks if no other requests are using them (reference count = 0). In this example, we free request 1 and block 2, 3, 4, 8 associated with it. We can see that the freed blocks are added to the tail of the free queue in the *reverse* order. This is because the last block of a request must hash more tokens and is less likely to be reused by other requests. As a result, it should be evicted first. -![Free queue after a request us freed](../assets/design/prefix_caching/free.png) +![Free queue after a request is freed](../assets/design/prefix_caching/free.png) ### Eviction (LRU) diff --git a/docs/examples/README.md b/docs/examples/README.md index 9d6126a65c4..a9a127a4d5d 100644 --- a/docs/examples/README.md +++ b/docs/examples/README.md @@ -9,9 +9,10 @@ vLLM's examples are organized into the following categories: - **[`features/`](../../examples/features)** – Demonstrations of individual vLLM features: automatic prefix caching, speculative decoding, LoRA, structured outputs, prompt embedding, pause/resume, batch invariance, KV events, data parallelism, and more. - **[`reasoning/`](../../examples/reasoning)** – Examples for reasoning with vLLM. - **[`tool_calling/`](../../examples/tool_calling)** – Examples for function/tool calling with vLLM. -- **[`applications/`](../../examples/applications)** – Application examples such as chatbots and RAG (Retrieval-Augmented Generation). +- **[`applications/`](../../examples/applications)** – Application examples such as simpler api server, chatbots and RAG (Retrieval-Augmented Generation). - **[`rl/`](../../examples/rl)** – Reinforcement learning examples. - **[`deployment/`](../../examples/deployment)** – Examples for deploying vLLM in production. - **[`ray_serving/`](../../examples/ray_serving)** – Scalable serving using Ray. -- **[`disaggregated/`](../../examples/disaggregated)** – Examples for disaggregated serving (separate prefill and decode), including various kv cache connectors (LMCache, Mooncake, FlexKV, P2P NCCL) and failure recovery. +- **[`disaggregated/`](../../examples/disaggregated)** – Examples for Disaggregated P/D (Prefill/Decoding) inference, including various kv cache connectors (LMCache, Mooncake, FlexKV, P2P NCCL) and failure recovery. +- **[`scale_out/`](../../examples/scale_out)** – Examples for Token In <> Token Out API Server. - **[`observability/`](../../examples/observability)** – Metrics, logging, tracing (OpenTelemetry), and dashboards (Grafana, Perses). diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md index 8352d2f20e0..578343096df 100644 --- a/docs/features/disagg_prefill.md +++ b/docs/features/disagg_prefill.md @@ -20,7 +20,7 @@ Two main reasons: Now supports 9 types of connectors: - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling. -- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. +- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup. - **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as: ```bash diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 39b826bfd56..8a48d11be73 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -74,12 +74,15 @@ vllm serve \ | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | +| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | ## Secondary Tiers Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields. +The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`. + ### Filesystem (FS) The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage. @@ -90,6 +93,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage | `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). | | `n_read_threads` | no | `16` | Read-priority I/O threads (load path). | | `n_write_threads` | no | `16` | Write-priority I/O threads (store path). | +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. | Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency. @@ -119,6 +123,45 @@ To enable KV cache sharing between multiple vLLM instances using the same `root_ PYTHONHASHSEED=0 vllm serve ... ``` +### Object Store (OBJ) + +The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object store through the NIXL OBJ backend. + +| Key | Required | Default | Notes | +| --- | --- | --- | --- | +| `type` | yes | — | Must be `obj`. | +| `store_config` | yes | — | Object store connection parameters (see below). | +| `prefix` | no | `""` | Key prefix prepended to all object keys. | +| `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. | +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. | + +`store_config` fields: + +| Key | Required | Default | Notes | +| --- | --- | --- | --- | +| `bucket` | yes | — | Bucket name. | +| `endpoint_override` | yes | — | Object store endpoint host; the URL scheme is set separately via `scheme`. | +| `scheme` | no | `http` | `http` or `https`. | +| `access_key`, `secret_key`, `session_token` | no | `""` | Explicit credentials. When left empty, the NIXL OBJ plugin falls back to the AWS SDK default credential provider chain (IAM roles, environment variables, credential files), which enables workload-identity auth on Kubernetes. | +| `region` | no | `""` | Bucket region, if the endpoint requires one. | +| `ca_bundle` | no | `""` | CA bundle path for TLS verification. | + +Object keys follow the same run-configuration digest scheme as the filesystem tier (see [On-Disk Layout](#on-disk-layout)) and are stored under the optional `prefix`. The [Cross-Process Sharing](#cross-process-sharing) requirement (`PYTHONHASHSEED`) applies to shared buckets as well, so instances sharing a bucket produce identical keys for identical content. At startup the tier probes object store connectivity and fails fast with a configuration error if the bucket is unreachable. + +### P2P (Including P/D) + +The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. + +| Key | Required | Default | Notes | +| --- | --- | --- | --- | +| `type` | yes | — | Must be `p2p`. | +| `host` | no | `0.0.0.0` | Address the control socket binds to. | +| `port` | no | `7777` | Port for the control socket. Must be reachable from peers. | +| `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. | +| `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. | + +The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process. + ## Tuning Tips - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload. @@ -127,6 +170,29 @@ PYTHONHASHSEED=0 vllm serve ... - FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high. - Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk. +## Per-Request Selective Offload + +Individual requests can cap how many of their tokens are eligible for offload by setting `max_offload_tokens` in the request's `kv_transfer_params`. Only the first `max_offload_tokens` tokens of the request are offloaded; blocks beyond that point are skipped on the store path. This is useful when a known prefix (e.g., a system prompt or shared context) is worth caching but later request-specific tokens are not. + +| Key | Type | Notes | +| --- | --- | --- | +| `max_offload_tokens` | non-negative `int` | Upper bound on tokens to offload for this request. `0` disables offload for the request entirely; omit the key (or set to `None`) for no cap. Non-`int`, negative, or `bool` values are rejected with a warning and treated as no cap. | + +!!! note + `max_offload_tokens` is experimental and subject to change. + +Example (OpenAI-compatible completions request): + +```json +{ + "model": "", + "prompt": "...", + "kv_transfer_params": { + "max_offload_tokens": 1024 + } +} +``` + ## Further Reading - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput). diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index f23acae10c4..cb857856b78 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -203,8 +203,10 @@ the vLLM JSON config. ### kv_connector_extra_config - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`. +- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`. - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`. - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`. +- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format). ## Notes diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 847743dfff1..df33cea0542 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -816,6 +816,44 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu export VLLM_VIDEO_FETCH_TIMEOUT= ``` +#### Video Decoding Backend + +vLLM decodes video bytes into frames using a selectable decoding backend. Three +backends are supported: + +- `opencv` (default): OpenCV-based decoder. +- `pyav`: PyAV decoder. +- `torchcodec`: TorchCodec (PyTorch-native) decoder. + +All three backends are ultimately backed by FFmpeg. `torchcodec` lets +you choose which FFmpeg version is used while `opencv` and `pyav` rely on +whichever FFmpeg build they were linked against. + +Select the backend by passing the `backend` parameter via `--media-io-kwargs`: + +```bash +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec"}}' +``` + +**TorchCodec-specific parameters:** + +The following parameters only apply to the `torchcodec` backend: + +- `num_ffmpeg_threads`: Number of FFmpeg decoding threads. `0` (default) relies + on the FFmpeg default, which is `min(cpu_count + 1, 16)`. This allows you to + control thread over-subscription. +- `seek_mode`: Seek mode for the decoder. `"exact"` (default) guarantees + frame-accurate sampling by scanning the file when the decoder is created. + `"approximate"` skips that scan for faster decoder creation, at the cost of + relying on the file's metadata (which may yield less accurate seeking). + +```bash +# Example: TorchCodec with approximate seek mode and 4 FFmpeg threads +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' +``` + #### Video Frame Recovery For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place. @@ -841,6 +879,55 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ Works with common video formats like MP4 when using OpenCV backends. +#### GPU Video Decoding with DeepStream (NVDEC) + +By default vLLM decodes video on the CPU. On NVIDIA GPUs you can instead decode +directly on the hardware video engine (NVDEC) with the DeepStream backend, which +keeps decoding off the CPU and can significantly increase video throughput. + +Install the backend (Linux x86-64 only): + +```bash +pip install vllm[deepstream] +``` + +The pip wheel bundles the DeepStream libraries but still relies on a few system +packages that pip cannot install. On Ubuntu: + +```bash +apt-get install -y \ + gstreamer1.0-tools gstreamer1.0-plugins-base gstreamer1.0-plugins-good \ + gstreamer1.0-plugins-bad gstreamer1.0-libav \ + python3-gi python3-gst-1.0 libv4l-0 cuda-libraries-13-0 +``` + +Select the backend either with an environment variable: + +```bash +export VLLM_VIDEO_LOADER_BACKEND=deepstream +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct +``` + +or per request via `--media-io-kwargs`: + +```bash +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "deepstream"}}' +``` + +**Parameters:** + +- `pool_size`: Number of GPU decode workers in the process-wide decode pool + (clamped to `[1, 16]`). When unset it defaults to + `VLLM_MEDIA_LOADING_THREAD_COUNT` (default `8`). The pool is a singleton, so + the first request's value wins. + +```bash +# Example: 12 decode workers +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "deepstream", "pool_size": 12}}' +``` + #### 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. diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 0f0cbd55354..03b05751c14 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -294,6 +294,21 @@ curl http://localhost:8000/v1/chat/completions \ !!! note The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine. +### Benchmarking the multi-turn proxy + +[`benchmarks/multi_turn/benchmark_serving_multi_turn.py`](../../benchmarks/multi_turn/benchmark_serving_multi_turn.py) supports targeting the disaggregated multi-turn proxy with the `--send-conversation-id` flag, which injects a per-conversation `conversation_id` into every request payload so the proxy can key cross-turn KV cache reuse. + +The flag is **off by default** so the benchmark is compatible with strict OpenAI-compatible frontends that reject unknown top-level fields. When benchmarking the multi-turn proxy you must pass it explicitly — otherwise every turn lands as a cache MISS and the bidirectional KV transfer path is never exercised. + +```bash +python benchmarks/multi_turn/benchmark_serving_multi_turn.py \ + --model --served-model-name \ + --url http://:8000 \ + --input-file benchmarks/multi_turn/generate_multi_turn.json \ + --num-clients 2 --max-active-conversations 6 \ + --send-conversation-id +``` + ### Limitations - Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns. @@ -408,6 +423,54 @@ To enable this feature: --kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}' ``` +## Metrics Reference + +vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer +activity for the last reporting interval. Example output: + +```text +KV Transfer metrics: Num successful transfers=4, Avg xfer time (ms)=1.381, +P90 xfer time (ms)=2.601, Avg post time (ms)=0.672, P90 post time (ms)=0.801, +Avg MB per transfer=2.25, Throughput (MB/s)=1629.549, Avg number of descriptors=72.0 +``` + +The table below describes each field. All timing values cover only the +successful transfers recorded in the current interval; failed transfers are +counted separately via Prometheus (see +[Prometheus metrics](#prometheus-metrics) below). + +| Metric | Unit | Description | +| -------- | ------ | ------------- | +| `Num successful transfers` | count | Number of NIXL KV-block transfers that completed without error during the interval. A transfer corresponds to one prefill request's worth of KV cache being moved from the prefiller to the decoder (or vice versa in bidirectional mode). | +| `Avg xfer time (ms)` | ms | Mean end-to-end transfer duration (`xferDuration` in NIXL telemetry, converted from µs). Measured from when the request is posted to when the backend reports completion, so it includes both the posting step and the actual data movement. | +| `P90 xfer time (ms)` | ms | 90th-percentile transfer duration. Use this to identify tail latency: a large gap between average and P90 suggests occasional stragglers (e.g., network congestion or large KV blocks). | +| `Avg post time (ms)` | ms | Mean time to submit the transfer request to the RDMA backend (`postDuration` in NIXL telemetry). This is the synchronous cost of posting work to the NIC queue (descriptor setup, etc.) before the async data movement begins. | +| `P90 post time (ms)` | ms | 90th-percentile request-posting duration. Elevated P90 here (with low xfer P90) points to overhead in submitting requests rather than in the data transfer itself. | +| `Avg MB per transfer` | MB | Mean payload size per transfer, computed as `total bytes transferred / number of transfers`. Reflects the average KV cache footprint of a single request (sequence length × layers × head dimension × dtype bytes). | +| `Throughput (MB/s)` | MB/s | Effective bandwidth over the interval: `total MB transferred / total xfer time (s)` across all successful transfers. This is aggregate throughput, not per-request bandwidth. | +| `Avg number of descriptors` | count | Mean number of NIXL memory descriptors (scatter-gather segments) submitted per transfer. More descriptors indicate more fragmented or larger KV cache allocations; very high counts can increase descriptor-registration overhead. | + +### Prometheus metrics + +In addition to the periodic log line, the following Prometheus metrics are +exported when NixlConnector is active: + +| Metric name | Type | Description | +| ------------- | ------ | ------------- | +| `vllm:nixl_xfer_time_seconds` | Histogram | Per-transfer RDMA copy duration (seconds). | +| `vllm:nixl_post_time_seconds` | Histogram | Time to submit the transfer request to the RDMA backend (seconds). | +| `vllm:nixl_bytes_transferred` | Histogram | Bytes moved per transfer. | +| `vllm:nixl_num_descriptors` | Histogram | Descriptor count per transfer. | +| `vllm:nixl_num_failed_transfers` | Counter | Cumulative count of failed NIXL KV-block transfers. | +| `vllm:nixl_num_failed_notifications` | Counter | Cumulative count of failed completion notifications (`send_notif`). | +| `vllm:nixl_num_kv_expired_reqs` | Counter | Requests whose KV blocks expired on the prefiller before the decoder read them (tracked on the P instance). | + +!!! tip + High `vllm:nixl_num_kv_expired_reqs` indicates that the prefiller's lease + duration (`kv_lease_duration`) is too short for your network or workload. + Increase it via `--kv-transfer-config '{"kv_connector_extra_config": + {"kv_lease_duration": }}'`. + ## Example Scripts/Code Refer to these example scripts in the vLLM repository: diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md new file mode 100644 index 00000000000..9bc64d2b86d --- /dev/null +++ b/docs/features/per_request_metrics.md @@ -0,0 +1,127 @@ +# Per-Request Metrics + +vLLM can return per-request timing metrics directly in API responses. +This is useful for billing, SLA monitoring, and latency analysis at the +individual request level, as a complement to the server-aggregated Prometheus +metrics exposed at `/metrics`. + +## Enabling + +Start the server with `--enable-per-request-metrics`: + +```bash +vllm serve meta-llama/Llama-3.1-8B-Instruct --enable-per-request-metrics +``` + +When this flag is set, supported API responses include metrics for each +attributable request. + +!!! note + At high concurrency, enabling per-request metrics computation may introduce + non-negligible CPU overhead. Benchmark your specific workload to evaluate the + impact before enabling in production. + +## Response Format + +When per-request metrics are enabled, the response includes a `metrics` object: + +```json +{ + "id": "chatcmpl-abc123", + "object": "chat.completion", + "model": "meta-llama/Llama-3.1-8B-Instruct", + "choices": [ ... ], + "usage": { + "prompt_tokens": 42, + "completion_tokens": 128, + "total_tokens": 170 + }, + "metrics": { + "time_to_first_token_ms": 85.2, + "generation_time_ms": 1240.5, + "queue_time_ms": 12.3, + "mean_itl_ms": 9.1, + "tokens_per_second": 103.2 + } +} +``` + +| Field | Description | +| --- | --- | +| `time_to_first_token_ms` | Time from when the request was scheduled until the first output token was generated (TTFT). | +| `generation_time_ms` | Decode time: time from the first output token to the last output token. Excludes both queue wait and prefill/TTFT. | +| `queue_time_ms` | Time the request spent waiting in the scheduler queue before processing began. | +| `mean_itl_ms` | Mean inter-token latency (average time between successive output tokens) during the decode phase. `null` for single-token responses. | +| `tokens_per_second` | Overall output token throughput: all generated tokens over the inference interval (scheduling to last output token). Unlike `generation_time_ms`, this includes the prefill phase, so it reflects end-to-end generation speed rather than pure decode speed. | + +All fields are `null` if the underlying timing data is not available for that +request. + +!!! note + Timing metrics describe a single generation stream, so they are only + returned when the request maps to exactly one. They are suppressed (the + `metrics` object is `null`) for requests with `n > 1`, because the + underlying timing data reflects only one of the `n` sequences and cannot be + accurately attributed to the request as a whole. Token usage + (`prompt_tokens`, `completion_tokens`) remains accurate in these cases. + Per-request metrics also require server-side statistics logging, which is + on by default. vLLM rejects `--enable-per-request-metrics` when + `--disable-log-stats` is also set. + +## Example Request + +=== "Non-streaming" + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + response = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + ) + + print(response.usage) + print(response.model_extra.get("metrics")) + ``` + +=== "Streaming" + + In streaming responses, metrics are attached to the final usage chunk (the + chunk sent after all content chunks). That chunk is only emitted when usage + reporting is enabled with `stream_options.include_usage: true` or forced + server-side with `--enable-force-include-usage`. Without forced usage, a + streaming client must set `stream_options.include_usage: true` to receive + metrics. + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + stream = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + stream=True, + stream_options={"include_usage": True}, + ) + + for chunk in stream: + if chunk.usage: + print("Usage:", chunk.usage) + print("Metrics:", chunk.model_extra.get("metrics")) + ``` + +## Completions API + +Per-request metrics are also available on the `/v1/completions` endpoint using +the same `metrics` response field. As with `n > 1`, metrics are omitted for +requests with multiple prompts, because the timing data cannot be attributed to +a single prompt's generation. + +## Relationship to Prometheus Metrics + +The `metrics` response field provides per-request values for a single request. +The `/metrics` Prometheus endpoint exposes server-level histograms (e.g. +`vllm:time_to_first_token_seconds`) that aggregate across all requests. diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 2be357d8860..69ece360761 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -9,7 +9,6 @@ The following are the supported quantization formats for vLLM: - [AutoAWQ](auto_awq.md) - [BitsAndBytes](bnb.md) -- [GGUF](gguf.md) - [GPTQModel](gptqmodel.md) - [Intel Neural Compressor](inc.md) - [LLM Compressor](llm_compressor/README.md) diff --git a/docs/features/quantization/auto_awq.md b/docs/features/quantization/auto_awq.md index e93005f2632..39dfd6fec11 100644 --- a/docs/features/quantization/auto_awq.md +++ b/docs/features/quantization/auto_awq.md @@ -49,7 +49,7 @@ To run an AWQ model with vLLM, you can use [TheBloke/Llama-2-7b-Chat-AWQ](https: ```bash python examples/deployment/llm_engine_example.py \ --model TheBloke/Llama-2-7b-Chat-AWQ \ - --quantization awq + --quantization auto_awq ``` AWQ models are also supported directly through the LLM entrypoint: @@ -70,7 +70,7 @@ AWQ models are also supported directly through the LLM entrypoint: sampling_params = SamplingParams(temperature=0.8, top_p=0.95) # Create an LLM. - llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="AWQ") + llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="auto_awq") # Generate texts from the prompts. The output is a list of RequestOutput objects # that contain the prompt, generated text, and other information. outputs = llm.generate(prompts, sampling_params) diff --git a/docs/features/quantization/gguf.md b/docs/features/quantization/gguf.md index 41912a50601..0aa76d679e1 100644 --- a/docs/features/quantization/gguf.md +++ b/docs/features/quantization/gguf.md @@ -3,8 +3,14 @@ !!! warning Please note that GGUF support in vLLM is highly experimental and under-optimized at the moment, it might be incompatible with other features. Currently, you can use GGUF as a way to reduce memory footprint. If you encounter any issues, please report them to the vLLM team. -!!! warning - Currently, vllm only supports loading single-file GGUF models. If you have a multi-files GGUF model, you can use [gguf-split](https://github.com/ggerganov/llama.cpp/pull/6135) tool to merge them to a single-file model. +!!! note + GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model. + +Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin): + +```bash +uv pip install vllm-gguf-plugin +``` To run a GGUF model with vLLM, you can use the `repo_id:quant_type` format to load directly from HuggingFace. For example, to load a Q4_K_M quantized model from [unsloth/Qwen3-0.6B-GGUF](https://huggingface.co/unsloth/Qwen3-0.6B-GGUF): diff --git a/docs/features/quantization/gptqmodel.md b/docs/features/quantization/gptqmodel.md index 636a952b655..235afee5f32 100644 --- a/docs/features/quantization/gptqmodel.md +++ b/docs/features/quantization/gptqmodel.md @@ -55,7 +55,7 @@ Here is an example of how to quantize `meta-llama/Llama-3.2-1B-Instruct`: ## Running a quantized model with vLLM -To run an GPTQModel quantized model with vLLM, you can use [DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2](https://huggingface.co/ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2) with the following command: +To run a GPTQModel quantized model with vLLM, you can use [DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2](https://huggingface.co/ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2) with the following command: ```bash python examples/deployment/llm_engine_example.py \ diff --git a/docs/features/quantization/inc.md b/docs/features/quantization/inc.md index adb6b3ae8e2..ffb90cec8c1 100644 --- a/docs/features/quantization/inc.md +++ b/docs/features/quantization/inc.md @@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \ --max-model-len 4096 ``` -!!! note - To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now. - ## Evaluating the Quantized Model with vLLM ```bash lm_eval --model vllm \ - --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \ + --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \ --tasks gsm8k \ --num_fewshot 5 \ --batch_size 128 diff --git a/docs/features/quantization/llm_compressor/int8_w8a8.md b/docs/features/quantization/llm_compressor/int8_w8a8.md index 21ed00d1393..64bce832c18 100644 --- a/docs/features/quantization/llm_compressor/int8_w8a8.md +++ b/docs/features/quantization/llm_compressor/int8_w8a8.md @@ -78,8 +78,6 @@ def tokenize(sample): ds = ds.map(tokenize, remove_columns=ds.column_names) ``` - - ### 3. Applying Quantization Now, apply the quantization algorithms: diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index 4b9571b38b9..a4da52557af 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -62,6 +62,8 @@ weight name. Unset fields fall back to the `--quantization` shorthand's defaults, or for already-quantized checkpoints to whatever the checkpoint declares. +On XPU, non-block FP8 scaled-mm linear layers default to W8A16; setting `--linear-backend xpu` forces W8A8. Use `--linear-backend xpu_woq` to explicitly select weight-only quantization (W8A16). + The CLI accepts the same shape as JSON or as dotted keys: ```bash diff --git a/docs/features/quantization/quantized_kvcache.md b/docs/features/quantization/quantized_kvcache.md index 2c5bfd64394..50b1c5c2df3 100644 --- a/docs/features/quantization/quantized_kvcache.md +++ b/docs/features/quantization/quantized_kvcache.md @@ -49,6 +49,32 @@ You can configure how the quantization scales are computed in vLLM using three d - `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs) - `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+ +### Skipping Specific Layers from KV-Cache Quantization + +Some attention layer types (e.g. sliding-window) are more sensitive to KV-cache quantization. The `--kv-cache-dtype-skip-layers` flag leaves the specified layers at the model's native dtype while keeping the rest of the layers under the chosen quantized dtype. The flag accepts either layer indices or layer-type names: + +```bash +# Skip every sliding-window attention layer. +vllm serve \ + --kv-cache-dtype fp8 \ + --kv-cache-dtype-skip-layers sliding_window + +# Skip specific layer indices. +vllm serve \ + --kv-cache-dtype fp8 \ + --kv-cache-dtype-skip-layers 0 1 23 +``` + +Programmatic usage: + +```python +llm = LLM( + model="meta-llama/Llama-3.1-8B-Instruct", + kv_cache_dtype="fp8", + kv_cache_dtype_skip_layers=["sliding_window"], +) +``` + --- ## Examples diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 92563a8b4bb..9495c4672ce 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -351,6 +351,69 @@ print(response.choices[0].message.reasoning) print(response.choices[0].message.content) ``` +## Suppressing Reasoning Output + +You can suppress reasoning content from API responses using the `include_reasoning` parameter. When set to `false`, reasoning tokens are still generated (so model quality is unaffected) but excluded from the response. This reduces network traffic without changing inference behavior. + +The parameter is supported in both the Chat Completions API and the Responses API, for streaming and non-streaming requests. + +When `include_reasoning=false`, vLLM also suppresses per-token metadata (logprobs and token IDs) to prevent leaking reasoning content through decoded token text in logprob entries or raw token IDs. + +### Chat Completions API + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") +model = client.models.list().data[0].id + +# Reasoning is included by default (include_reasoning=True) +response = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "What is 15 * 37?"}], + extra_body={"include_reasoning": False}, +) + +msg = response.choices[0].message +assert msg.content # Content is still present +assert not getattr(msg, "reasoning", None) # Reasoning is suppressed +``` + +Streaming works the same way, reasoning deltas are omitted from chunks: + +```python +stream = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "What is 15 * 37?"}], + stream=True, + extra_body={"include_reasoning": False}, +) + +for chunk in stream: + delta = chunk.choices[0].delta + # delta.reasoning will always be None + if delta.content: + print(delta.content, end="", flush=True) +``` + +### Responses API + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") + +response = client.responses.create( + model=client.models.list().data[0].id, + input="What is 15 * 37?", + include_reasoning=False, +) + +# No "reasoning" items in output +types = [item.type for item in response.output] +assert "reasoning" not in types +``` + ## Limitations - The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`). @@ -439,7 +502,7 @@ Additionally, to enable structured output, you'll need to create a new `Reasoner end_token: str = "" @classmethod - def from_tokenizer(cls, tokenizer: PreTrainedTokenizer) -> Reasoner: + def from_tokenizer(cls, tokenizer: PythonBackend) -> Reasoner: return cls( start_token_id=tokenizer.encode("", add_special_tokens=False)[0], end_token_id=tokenizer.encode("", add_special_tokens=False)[0], diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 58d1df9dced..ceb25dbfd02 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -17,6 +17,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods - [Suffix Decoding](suffix.md) - [Hidden State Extraction](extract_hidden_states.md) - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental) +- [Dynamic Speculative Decoding](dynamic_speculative_decoding.md) ## Method Selection at a Glance @@ -33,6 +34,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings. | N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. | | Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. | | Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). | +| Dynamic Speculative Decoding | High gain | Higher than base SD method | Useful for RL or workload with fluctuating QPS | For reproducible measurements in your environment, use [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py) @@ -84,6 +86,7 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and | `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]`. | + | `use_heterogeneous_vocab` | `boolean` | `false` | Allow draft and target models with different vocabularies. Builds a token-level intersection at initialisation and constrains draft logits to shared tokens only. Only compatible with `method=draft_model`. Probabilistic draft sampling (`draft_sample_method='probabilistic'`) is not yet supported when this option is enabled. | !!! note Gemma 4 assistant checkpoints are handled as Gemma 4 MTP speculators, not @@ -140,6 +143,33 @@ vllm serve \ }' ``` +#### Cross-Vocabulary Draft Models (TLI) + + By default, vLLM requires the draft and target models to share the same + vocabulary. Setting `use_heterogeneous_vocab: true` enables the + **Token-Level Intersection (TLI)** algorithm, which allows draft models + from a different model family with a different tokenizer. + + At initialisation, vLLM builds a mapping between the two vocabularies by + normalising token strings and computing their intersection. Draft logits are + constrained to the shared tokens before sampling, and the sampled token IDs + are translated to the target vocabulary before rejection sampling. + + ```python + from vllm import LLM, SamplingParams + + llm = LLM( + model="Qwen/Qwen3-8B", + speculative_config={ + "method": "draft_model", + "model": "HuggingFaceTB/SmolLM2-135M-Instruct", + "num_speculative_tokens": 3, + "use_heterogeneous_vocab": True, + }, + gpu_memory_utilization=0.5, + ) +``` + ### Notes - `--speculative-config` expects a JSON object on the CLI. In YAML config @@ -151,6 +181,7 @@ vllm serve \ - 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. +- `use_heterogeneous_vocab` currently supports greedy draft sampling only. Probabilistic acceptance (temperature > 0 draft sampling) is not yet supported and will be added in a future release. ## Lossless guarantees of Speculative Decoding @@ -188,7 +219,7 @@ For mitigation strategies, please refer to the FAQ entry *Can the output of a pr ## Known Feature Incompatibility -1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0` +1. Pipeline parallelism is not composable with speculative decoding as of `vllm<=0.15.0` 2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0` ## Resources for vLLM contributors diff --git a/docs/features/speculative_decoding/draft_model.md b/docs/features/speculative_decoding/draft_model.md index b4662e6438f..636c797324c 100644 --- a/docs/features/speculative_decoding/draft_model.md +++ b/docs/features/speculative_decoding/draft_model.md @@ -76,6 +76,34 @@ The code used to request as completions as a client remains unchanged: print(completion) ``` +## Draft Model Method with heterogeneous vocabs + + By default, vLLM requires the draft and target models to share the same vocabulary. Setting `use_heterogeneous_vocab: true` enables the **Token-Level Intersection (TLI)** algorithm, which allows draft models from a different model family with a different tokenizer. + + Currently,`use_heterogeneous_vocab` currently requires `draft_sample_method='greedy'` (the default). Probabilistic draft sampling is not yet supported and will be added in a + future release. + + ```python + from vllm import LLM, SamplingParams + + llm = LLM( + model="Qwen/Qwen3-8B", + speculative_config={ + "method": "draft_model", + "model": "HuggingFaceTB/SmolLM2-135M-Instruct", + "num_speculative_tokens": 3, + "use_heterogeneous_vocab": True, + }, + gpu_memory_utilization=0.5, + ) +outputs = llm.generate(prompts，sampling_params) + +for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") +``` + !!! warning Note: Please use `--speculative-config` to set all configurations related to speculative decoding. The previous method of specifying the model diff --git a/docs/features/speculative_decoding/dynamic_speculative_decoding.md b/docs/features/speculative_decoding/dynamic_speculative_decoding.md new file mode 100644 index 00000000000..682eaafd29d --- /dev/null +++ b/docs/features/speculative_decoding/dynamic_speculative_decoding.md @@ -0,0 +1,76 @@ +# Dynamic Speculative Decoding + +## Why is Dynamic SD needed? + +SD methods need to verify K tokens for each sequence during decoding. As BS increases, the effective BS becomes BS\*K which increases the compute requirement during verification. When this BS\*K goes beyond a critical BS then SD negatively impacts the decode speed (TPOT). DSD helps by tuning the K to an optimal value such that we continue to reap the benefits from SD. + +## Use cases + +* Variable concurrency workload using same deployment. K would decrease as concurrency increases. +* During RL rollout where we start off with high BS but then end up with small BS due to very few long tail request which end up generating a lot of tokens stalling the progress of the current rollout. Here K would go up during the end of rollout. + +## `--speculative-config` schema + +To use Dynamic SD, add `num_speculative_tokens_per_batch_size` to the config of an SD method which is a list of list. Here, an entry is `[start_bs, end_bs, optimal_K]` which means when the concurrency is within range `[start_bs, end_bs]` then `optimal_K` number of draft tokens are used. For e.g., + +```bash +--speculative-config '{ + "method": "eagle", + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' +``` + +implies that: + +* K=3 will be used when the concurrency is in range [1, 64] +* K=1 will be used when the concurrency is in range [65, 128] +* K=0 will be used when the concurrency is in range [129, 512], i.e., no draft tokens will be produced. + +## Online Examples + +### Dynamic SD Eagle Drafter + +```bash +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --speculative-config '{ + "method": "eagle", + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' +``` + +### Dynamic SD Eagle3 Drafter + +```bash +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --speculative-config '{ + "method": "eagle3", + "model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 16, 5], + [17, 32, 4], + [33, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' + +``` + +## Limitations + +* Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box +* Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature +* Not compatible with data parallelism (`--data-parallel-size > 1`). Each DP rank schedules independently, so ranks can pick different K values, causing DP collective divergence and deadlocks. When DP is enabled, vLLM automatically disables `num_speculative_tokens_per_batch_size` and falls back to the static `num_speculative_tokens` value. diff --git a/docs/features/speculative_decoding/extract_hidden_states.md b/docs/features/speculative_decoding/extract_hidden_states.md index 2184a71f489..b7df376d9ff 100644 --- a/docs/features/speculative_decoding/extract_hidden_states.md +++ b/docs/features/speculative_decoding/extract_hidden_states.md @@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import ( with tempfile.TemporaryDirectory() as tmpdir: llm = LLM( model="Qwen/Qwen3-8B", - enable_chunked_prefill=False, speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1, @@ -59,17 +58,58 @@ For improved performance, it is recommended to use a RAM-mounted file system suc ```bash vllm serve Qwen/Qwen3-8B \ --speculative_config '{"method": "extract_hidden_states", "num_speculative_tokens": 1, "draft_model_config": {"hf_config": {"eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4]}}}' \ - --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' \ - --no-enable-chunked-prefill + --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' +``` + +## Per-Request Options + +Both offline and online modes support per-request options via `kv_transfer_params`: + +| Parameter | Default | Description | +| --- | --- | --- | +| `hidden_states_path` | Auto-generated | Custom file path for saving hidden states. If not set, files are saved to `/.safetensors`. Requires `allow_custom_save_path` to be enabled in the server config. | +| `include_output_tokens` | `False` | When `True`, save hidden states for both prompt and generated output tokens. When `False`, only prompt token hidden states are saved. | + +### Offline usage + +Pass per-request options via `extra_args` on `SamplingParams`: + +```python +SamplingParams( + max_tokens=32, + extra_args={ + "kv_transfer_params": { + "hidden_states_path": "/tmp/my_output.safetensors", + "include_output_tokens": True, + } + }, +) +``` + +### Online usage + +Pass `kv_transfer_params` as a top-level field in the API request: + +```json +{ + "model": "Qwen/Qwen3-8B", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 32, + "kv_transfer_params": { + "hidden_states_path": "/tmp/my_output.safetensors", + "include_output_tokens": true + } +} ``` ## Configuration -The `kv_connector_extra_config` dict accepts these options: +The `kv_connector_extra_config` dict accepts these server-level options: | Parameter | Default | Description | | --- | --- | --- | -| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved | +| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved (used when `hidden_states_path` is not set per-request) | +| `allow_custom_save_path` | `False` | Allow API clients to specify custom file paths via `hidden_states_path`. When disabled, client-provided paths are ignored with a warning. Enable only with trusted clients — custom paths can write to arbitrary locations on the server. | | `num_writer_threads` | `8` | Thread pool size for async disk writes | | `use_synchronization_lock` | `True` | Use file locks so concurrent readers block until writes complete. Can be disabled for batch generation where synchronization is not needed. | diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 95092734f3d..ae65231919a 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -109,24 +109,30 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t ## Constrained Decoding Behavior -Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode: +Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode and the per-tool `strict` field: | `tool_choice` value | Schema-constrained decoding | Behavior | | --- | --- | --- | | Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. | | `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. | -| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. | +| `"auto"` | Only when `strict: true` is set on at least one tool | Structural-tag parsers constrain tool-call arguments when a tool opts in with `strict: true`. Without it, the model generates freely and tool calls are extracted from raw text. | | `"none"` | N/A | No tool calls are produced. | -When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`. +### Strict Mode -### Strict Mode (`strict` parameter) +For `tool_choice="required"` or named function calling, structural-tag constraints are always applied regardless of the `strict` field. For `tool_choice="auto"`, setting `strict: true` on at least one tool opts in to structural-tag constraints; without it, the model generates freely and tool calls are extracted from raw text. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages. -The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode. +For best compatibility with strict schema enforcement, define tool parameter schemas in the OpenAI strict-schema style: -vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic. +* Set `additionalProperties` to `false` for each object in `parameters`. +* Mark all fields in `properties` as required. +* Represent optional fields by allowing `null`, for example `{"type": ["string", "null"]}`. -Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313). +vLLM also provides a global toggle via the `VLLM_ENFORCE_STRICT_TOOL_CALLING` environment variable (defaults to `true`). When set to `false`, vLLM does not attach structural tags for tool calling regardless of the per-tool `strict` field. This environment variable only affects structural-tag based tool calling; it does not change schema-derived structured outputs used by named function calling or `tool_choice="required"`. + +```bash +VLLM_ENFORCE_STRICT_TOOL_CALLING=false vllm serve ... +``` ## Automatic Function Calling @@ -146,7 +152,7 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template! !!! note - With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details. + With `tool_choice="auto"`, schema-level constraint requires both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`. When these conditions are met and the selected parser supports structural tags, vLLM constrains tool-call arguments. Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema. ### Hermes Models (`hermes`) @@ -315,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup Flags: `--tool-call-parser hermes` -### MiniMax Models (`minimax_m1`) - -Supported models: - -* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja)) -* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja)) - -Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja` - ### DeepSeek-V3 Models (`deepseek_v3`) Supported models: @@ -341,7 +338,7 @@ Supported models: Flags: `--tool-call-parser deepseek_v31 --chat-template {see_above}` -### OpenAI OSS Models ('openai`) +### OpenAI OSS Models (`openai`) Supported models: @@ -504,6 +501,13 @@ Flags: `--tool-call-parser pythonic --chat-template {see_above}` !!! warning Llama's smaller models frequently fail to emit tool calls in the correct format. Results may vary depending on the model. +## Benchmarking Tool-Calling Performance + +To measure serving latency and throughput on realistic tool-calling traffic, +use the BFCL (Berkeley Function Calling Leaderboard) dataset with +`vllm bench serve`. See the [BFCL benchmark example](../benchmarking/cli.md#bfcl-tool-calling-benchmark) +for the full server + client commands. + ## How to Write a Tool Parser Plugin A tool parser plugin is a Python file containing one or more ToolParser implementations. You can write a ToolParser similar to the `Hermes2ProToolParser` in [vllm/tool_parsers/hermes_tool_parser.py](../../vllm/tool_parsers/hermes_tool_parser.py). diff --git a/docs/getting_started/installation/cpu.apple.inc.md b/docs/getting_started/installation/cpu.apple.inc.md index e54afc49384..479b6d2c011 100644 --- a/docs/getting_started/installation/cpu.apple.inc.md +++ b/docs/getting_started/installation/cpu.apple.inc.md @@ -15,6 +15,10 @@ Currently the CPU implementation for macOS supports FP32 and FP16 datatypes. - SDK: `XCode 15.4` or later with Command Line Tools - Compiler: `Apple Clang >= 15.0.0` +!!! note + The macOS CPU build is smoke-tested in CI on the latest GA Apple Silicon + runner; other macOS or Apple Clang versions are best-effort. + --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -31,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match +uv pip install -r requirements/cpu.txt uv pip install -e . ``` -!!! tip - The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts. - - The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation. - !!! note On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device. diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index f01ba429ee0..3950adc0251 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -20,12 +20,12 @@ Pre-built vLLM wheels for Arm are available since version 0.11.2. These wheels c ```bash export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl --torch-backend cpu +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --torch-backend cpu ``` ??? console "pip" ```bash - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" @@ -63,7 +63,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/nightly/cpu --index If you insist on using `pip`, you have to specify the full URL (link address) of the wheel file (which can be obtained from https://wheels.vllm.ai/nightly/cpu/vllm). ```bash - pip install https://wheels.vllm.ai/4fa7ce46f31cbd97b4651694caf9991cc395a259/vllm-0.13.0rc2.dev104%2Bg4fa7ce46f.cpu-cp38-abi3-manylinux_2_35_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu # current nightly build (the filename will change!) + pip install https://wheels.vllm.ai/2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7/vllm-0.23.1rc1.dev901%2Bg2f3f441f8.cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu # current nightly build (the filename will change!) ``` #### Install specific revisions @@ -96,8 +96,8 @@ cd vllm_source Third, install required dependencies: ```bash -uv pip install -r requirements/build/cpu.txt --torch-backend cpu -uv pip install -r requirements/cpu.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu --index-strategy unsafe-best-match +uv pip install -r requirements/cpu.txt --torch-backend cpu --index-strategy unsafe-best-match ``` ??? console "pip" diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 7225d1d6c77..8b3605e8557 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -142,6 +142,10 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu === "IBM Z (S390X)" --8<-- "docs/getting_started/installation/cpu.s390x.inc.md:build-image-from-source" +## AMD Zen optimizations {#amd-zen-optimizations} + +--8<-- "docs/getting_started/installation/cpu.x86.inc.md:amd-zen-optimizations" + ## Related runtime environment variables - `VLLM_CPU_KVCACHE_SPACE`: specify the KV Cache size (e.g, `VLLM_CPU_KVCACHE_SPACE=40` means 40 GiB space for KV cache), larger setting will allow vLLM to run more requests in parallel. This parameter should be set based on the hardware configuration and memory management pattern of users. Default value is `0`. @@ -149,12 +153,14 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu - `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`. - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence. - `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False). +- `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations). ## FAQ ### Which `dtype` should be used? - Currently, vLLM CPU uses model default settings as `dtype`. However, due to unstable float16 support in torch CPU, it is recommended to explicitly set `dtype=bfloat16` if there are any performance or accuracy problem. +- On AMD Zen CPUs (`ZenCpuPlatform`), `float16` is **not** supported. Only `bfloat16` and `float32` are accepted; models declared with `float16` are auto-downcast to `bfloat16` at model load time. See [AMD Zen optimizations](#amd-zen-optimizations). ### How to launch a vLLM service on CPU? @@ -227,6 +233,25 @@ By providing MODEL_FILTER and DTYPE_FILTER, only commands for related model ID a ON_CPU=1 SERVING_JSON=serving-tests-cpu-text.json DRY_RUN=1 MODEL_FILTER=meta-llama/Llama-3.1-8B-Instruct DTYPE_FILTER=bfloat16 bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh ``` +### How do I enable AMD Zen optimizations? {#how-do-i-enable-amd-zen-optimizations} + +On an AMD Zen 4 / Zen 5 CPU, install the CPU wheel with the `zen` extra so vLLM pulls the tested `zentorch` version for that release: + +```bash +export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') +uv pip install "vllm[zen]" --extra-index-url https://wheels.vllm.ai/${VLLM_VERSION}/cpu --index-strategy first-index --torch-backend cpu +``` + +vLLM auto-detects the platform and routes linear layers through ZenDNN-optimized kernels - no flag needed. To verify it is engaged, look for the platform-selection line in the server's startup logs: + +```bash +vllm serve Qwen/Qwen3-0.6B 2>&1 | grep "AMD Zen CPU detected with zentorch installed" +``` + +For per-backend dispatch details (which kernel each linear layer was bound to), re-run with `VLLM_LOGGING_LEVEL=DEBUG` and grep for `CPU unquantized GEMM dispatch`. + +See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, supported dtypes, and the `VLLM_ZENTORCH_WEIGHT_PREPACK` knob. + ### How to decide `VLLM_CPU_OMP_THREADS_BIND`? - Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and 1 CPU per rank will be reserved for other vLLM components when `world_size > 1`. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following. diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 1e36b431764..15baa487c2a 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source. ```bash uv pip install -v \ - --extra-index-url https://download.pytorch.org/whl/cpu \ - --torch-backend auto \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ + --torch-backend cpu \ + --index-strategy unsafe-best-match && \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv 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 ad051d22dc8..6ded3b50832 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -1,4 +1,4 @@ - + --8<-- [start:installation] vLLM supports basic model inferencing and serving on x86 CPU platform, with data types FP32, FP16 and BF16. @@ -24,13 +24,13 @@ Pre-built vLLM wheels for x86 with AVX512/AVX2 are available since version 0.17. export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') # use uv -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl --torch-backend cpu +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --torch-backend cpu ``` ??? console "pip" ```bash # use pip - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" Before use vLLM CPU installed via wheels, make sure TCMalloc and Intel OpenMP are installed and added to `LD_PRELOAD`: @@ -88,8 +88,8 @@ cd vllm_source Install the required dependencies: ```bash -uv pip install -r requirements/build/cpu.txt --torch-backend cpu -uv pip install -r requirements/cpu.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu --index-strategy unsafe-best-match +uv pip install -r requirements/cpu.txt --torch-backend cpu --index-strategy unsafe-best-match ``` ??? console "pip" @@ -200,7 +200,19 @@ docker build -f docker/Dockerfile.cpu \ --target vllm-openai . ``` -#### Launching the OpenAI server +#### Building with AMD Zen optimizations + +For AMD Zen 4 / Zen 5 hosts (`linux/amd64` only), use the `vllm-openai-zen` target. It extends the default `vllm-openai` image and adds `zentorch` via the `vllm[zen]` extra so `ZenCpuPlatform` auto-activates at runtime: + +```bash +docker build -f docker/Dockerfile.cpu \ + --tag vllm-cpu-zen-env \ + --target vllm-openai-zen . +``` + +The resulting image accepts the same arguments and environment variables as `vllm-openai` (see [Launching the OpenAI server](#launching-the-openai-server) below); no extra flag is needed to engage Zen optimizations. See [AMD Zen optimizations](cpu.md#amd-zen-optimizations) for runtime behavior and the supported-dtype caveats. + +#### Launching the OpenAI server {#launching-the-openai-server} ```bash docker run --rm \ @@ -216,5 +228,36 @@ docker run --rm \ ``` --8<-- [end:build-image-from-source] +--8<-- [start:amd-zen-optimizations] + +On AMD Zen CPUs, vLLM auto-selects `ZenCpuPlatform` (a subclass of `CpuPlatform`) which dispatches linear layers through [`zentorch`](https://github.com/amd/ZenDNN-pytorch-plugin)'s ZenDNN-optimized kernels. See the FAQ entry [How do I enable AMD Zen optimizations?](#how-do-i-enable-amd-zen-optimizations) for the install command. + +### Detection rules + +`ZenCpuPlatform` is selected when **all** of the following hold: + +- vLLM is built for CPU +- `/proc/cpuinfo` reports `AuthenticAMD` and `avx512` +- `import zentorch` succeeds + +Otherwise, vLLM falls back to the default `CpuPlatform` (oneDNN / sgl-kernel paths). + +### Supported dtypes + +`float16` is **not** supported on `ZenCpuPlatform`. `ZenCpuPlatform.supported_dtypes` advertises only `bfloat16` and `float32`, so models declared with `torch_dtype=float16` are auto-downcast to `bfloat16` at load time with the standard `"Your device 'cpu' doesn't support torch.float16. Falling back to torch.bfloat16 for compatibility."` warning emitted from `vllm/config/model.py`. + +### Environment variables + +- `VLLM_ZENTORCH_WEIGHT_PREPACK` (default `1`): eagerly prepacks linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Set to `0` to disable. + +### Docker + +The `vllm-openai-zen` Docker target (in `docker/Dockerfile.cpu`) extends the default `vllm-openai` image with `vllm[zen]`. Build it with `docker build -f docker/Dockerfile.cpu --target vllm-openai-zen .` — see [Building with AMD Zen optimizations](#building-with-amd-zen-optimizations) for the full command and run instructions. + +### Reference + +For the design rationale, see [RFC #35089: In-Tree AMD Zen CPU Backend via zentorch](https://github.com/vllm-project/vllm/issues/35089). + +--8<-- [end:amd-zen-optimizations] --8<-- [start:extra-information] --8<-- [end:extra-information] diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index ec333b3ee1b..0e86c0e6049 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -43,7 +43,7 @@ As of now, vLLM's binaries are compiled with CUDA 12.9 and public PyTorch releas export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') export CUDA_VERSION=130 # or other export CPU_ARCH=$(uname -m) # x86_64 or aarch64 -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_35_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION} +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION} ``` #### Install the latest code @@ -68,8 +68,8 @@ uv pip install -U vllm \ If you insist on using `pip`, you have to specify the full URL of the wheel file (which can be obtained from the web page). ```bash - pip install -U https://wheels.vllm.ai/nightly/vllm-0.11.2.dev399%2Bg3c7461c18-cp38-abi3-manylinux_2_31_x86_64.whl # current nightly build (the filename will change!) - pip install -U https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-0.11.2.dev399%2Bg3c7461c18-cp38-abi3-manylinux_2_31_x86_64.whl # from specific commit + pip install -U https://wheels.vllm.ai/2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7/vllm-0.23.1rc1.dev901%2Bg2f3f441f8-cp38-abi3-manylinux_2_28_x86_64.whl # current nightly build (the filename will change!) + pip install -U https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-0.23.1rc1.dev901%2Bg2f3f441f8-cp38-abi3-manylinux_2_28_x86_64.whl # from specific commit ``` ##### Install specific revisions @@ -139,6 +139,15 @@ You can find more information about vLLM's wheels in [Install the latest code](# #### Full build (with compilation) {#full-build} +!!! note "Compiler requirement" + Building from source requires GCC/G++ ≥ 11.3. PyTorch's C++20 headers are + not compatible with GCC 10 or GCC < 11.3. On Ubuntu 22.04: + ```bash + sudo apt-get install -y gcc-11 g++-11 + sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 \ + --slave /usr/bin/g++ g++ /usr/bin/g++-11 + ``` + If you want to modify C++ or CUDA code, you'll need to build vLLM from source. This can take several minutes: ```bash diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index f8385997eea..59c9723e666 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -27,6 +27,19 @@ If you need a different ROCm version or want to use an existing PyTorch installa --8<-- [end:set-up-using-python] --8<-- [start:pre-built-wheels] +!!! warning "Python 3.12 required for ROCm wheels" + + ROCm pre-built wheels are only available for **Python 3.12**. If you are using a different Python version (e.g. 3.11 or 3.13), the installer **will silently fall back** to the CUDA wheel from PyPI, which will fail on AMD GPUs with errors like `libcudart.so: cannot open shared object file`. + + To check your Python version: `python3 --version` + + If you need Python 3.12, you can create an isolated environment with `uv`: + + ```bash + uv venv --python 3.12 --seed --managed-python + source .venv/bin/activate + ``` + To install the latest version of vLLM for Python 3.12, ROCm 7.0 and `glibc >= 2.35`. ```bash diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index f22f5159473..8564f2a7265 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels. - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers). - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)): +- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue. ```bash git clone https://github.com/vllm-project/vllm.git diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 194db05e395..07fbd7e4d55 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -32,7 +32,6 @@ def title(text: str) -> str: "mae": "MAE", "ner": "NER", "tpu": "TPU", - "gguf": "GGUF", "lora": "LoRA", "nccl": "NCCL", "rlhf": "RLHF", diff --git a/docs/models/hardware_supported_models/cpu.md b/docs/models/hardware_supported_models/cpu.md index 9c6dd9feb79..ddc519e8f16 100644 --- a/docs/models/hardware_supported_models/cpu.md +++ b/docs/models/hardware_supported_models/cpu.md @@ -1,5 +1,8 @@ # CPU - Intel® Xeon® +!!! note "AMD Zen CPUs" + On AMD Zen 4 / Zen 5 CPUs, AMD Zen optimizations are auto-enabled when the [`zentorch`](https://github.com/amd/ZenDNN-pytorch-plugin) package is installed. All models supported by vLLM on CPU are supported on AMD Zen as well; model compatibility does not change. This page reflects the current CPU reference validation matrix on Intel systems. See [AMD Zen optimizations](../../getting_started/installation/cpu.md#amd-zen-optimizations) for details. + ## Validated Hardware | Hardware | diff --git a/docs/models/hardware_supported_models/xpu.md b/docs/models/hardware_supported_models/xpu.md index cfda6c76f05..d065b4b6890 100644 --- a/docs/models/hardware_supported_models/xpu.md +++ b/docs/models/hardware_supported_models/xpu.md @@ -27,14 +27,12 @@ | Qwen/QwQ-32B | QwenForCausalLM | ✅ | | | | deepseek-ai/DeepSeek-V2-Lite | DeepSeekForCausalLM | ✅ | | | | meta-llama/Llama-3.1-8B-Instruct | LlamaForCausalLM | ✅ | | | -| baichuan-inc/Baichuan2-13B-Chat | BaichuanForCausalLM | ✅ | | | | THUDM/GLM-4-9B-chat | GLMForCausalLM | ✅ | | | | THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | | | 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | | | THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | | -| baichuan-inc/Baichuan2-13B-Chat | BaichuanForCausalLM | ✅ | | | | meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | | | THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | | diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index 2a5357e4fee..f8de9d437ad 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -2,8 +2,7 @@ !!! note We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance -improvements over using Hugging Face Transformers or Sentence Transformers directly. - + improvements over using Hugging Face Transformers or Sentence Transformers directly. We plan to optimize pooling models in vLLM. Please comment on if you have any suggestions! ## What are pooling models? @@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). !!! note 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. + are a subset of classification models that accept two prompts as input and output num_labels equal to 1. ### Pooling Types @@ -143,7 +142,7 @@ enabling the corresponding APIs. The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt. It is primarily designed for [classification models](classify.md). -For more information about `LLM.embed`, see [this page](classify.md#offline-inference). +For more information about `LLM.classify`, see [this page](classify.md#offline-inference). ### `LLM.embed` @@ -184,7 +183,7 @@ Our online Server provides endpoints that correspond to the offline APIs: - Corresponding to `LLM.classify`: - [Classification API](classify.md#online-serving)(`/classify`) - Corresponding to `LLM.score`: - - [Score API](scoring.md#score-api)(`/score`) + - [Score API](scoring.md#score-api) (`/score`, `/v1/score`) - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) - Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models. @@ -302,7 +301,7 @@ Pooling models now support token-wise task. ### Score task -`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels +`score` task has been removed in v0.21, use `classify` instead. 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. ### Pooling multitask support diff --git a/docs/models/pooling_models/classify.md b/docs/models/pooling_models/classify.md index 6860b09c31e..360f9294310 100644 --- a/docs/models/pooling_models/classify.md +++ b/docs/models/pooling_models/classify.md @@ -31,7 +31,6 @@ The most fundamental application of classification models is to categorize input | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | ------------------------------ | ------------------------------------------ | -| `ErnieForSequenceClassification` | BERT-like Chinese ERNIE | `Forrest20231206/ernie-3.0-base-zh-cls` | | | | `GPT2ForSequenceClassification` | GPT2 | `nie3e/sentiment-polish-gpt2-small` | | | | `Qwen2ForSequenceClassification`C | Qwen2-based | `jason9693/Qwen2.5-1.5B-apeach` | | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 47f85b7440e..1b9d14d7a0a 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -39,7 +39,6 @@ You can compute pairwise similarity scores to build a similarity matrix using th | ------------ | ------ | ----------------- | ------------------------------ | ------------------------------------------ | | `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | | `BertSpladeSparseEmbeddingModel` | SPLADE | `naver/splade-v3` | | | -| `ErnieModel` | BERT-like Chinese ERNIE | `shibing624/text2vec-base-chinese-sentence` | | | | `Gemma2Model`C | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | | `Gemma3TextModel`C | Gemma 3-based | `google/embeddinggemma-300m`, etc. | ✅︎ | ✅︎ | | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | diff --git a/docs/models/pooling_models/reward.md b/docs/models/pooling_models/reward.md index 4acacda5004..6049eb0a5f9 100644 --- a/docs/models/pooling_models/reward.md +++ b/docs/models/pooling_models/reward.md @@ -143,4 +143,4 @@ More examples can be found here: [examples/pooling/reward](../../../examples/poo ### `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. +`llm.reward` API is deprecated and was removed in v0.24. 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 c8b4c73cfb3..e3b54b02075 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -19,7 +19,7 @@ The score models is designed to compute similarity scores between two input prom - Offline APIs: - `LLM.score` - Online APIs: - - [Score API](scoring.md#score-api) (`/score`) + - [Score API](scoring.md#score-api) (`/score`, `/v1/score`) - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) !!! note @@ -157,7 +157,7 @@ A code example can be found here: [examples/basic/offline_inference/score.py](.. ### Score API -Our Score API (`/score`) is similar to `LLM.score`, compute similarity scores between two input prompts. +Our Score API (`/score`, `/v1/score`) is similar to `LLM.score`, compute similarity scores between two input prompts. #### Parameters @@ -440,7 +440,7 @@ More examples can be found here: [examples/pooling/score](../../../examples/pool ## Supported Features -AS cross-encoder models are a subset of classification models that accept two prompts as input and output num_labels equal to 1, cross-encoder features should be consistent with (sequence) classification. For more information, see [this page](classify.md#supported-features). +As cross-encoder models are a subset of classification models that accept two prompts as input and output num_labels equal to 1, cross-encoder features should be consistent with (sequence) classification. For more information, see [this page](classify.md#supported-features). ### Score Template diff --git a/docs/models/pooling_models/token_classify.md b/docs/models/pooling_models/token_classify.md index 5c4798935bf..6b2cefbde55 100644 --- a/docs/models/pooling_models/token_classify.md +++ b/docs/models/pooling_models/token_classify.md @@ -44,8 +44,8 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | --------------------------- | --------------------------------------- | | `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | | -| `ErnieForTokenClassification` | BERT-like Chinese ERNIE | `gyr66/Ernie-3.0-base-chinese-finetuned-ner` | | | | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | | +| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | | | `Qwen3ForTokenClassification`C | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 02050b7f50f..0c2a322e80f 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | | | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | | | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | | -| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | | +| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | | | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | | | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ | | `*ForConditionalGeneration`C, `*ForCausalLM`C, etc. | Generative models | \* | N/A | \* | \* | diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6fda22d1368..f19f01c6989 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang ### Transformers -vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend". +vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation. Currently, the Transformers modeling backend works for the following: @@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded: That's it! -For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class: +For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
configuration_my_model.py @@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported). + - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated. - `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s: - You only need to do this for layers which are not present on all pipeline stages - vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages + - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage - The `list` in the first element of the `tuple` contains the names of the input arguments - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code @@ -240,50 +242,24 @@ Use the Hugging Face CLI to [manage models](https://huggingface.co/docs/huggingf ```bash # List cached models -hf scan-cache +hf cache list -q # Show detailed (verbose) output -hf scan-cache -v +hf cache list # Specify a custom cache directory -hf scan-cache --dir ~/.cache/huggingface/hub +hf cache list --dir ~/.cache/huggingface/hub ``` #### Delete a cached model -Use the Hugging Face CLI to interactively [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache: +Use the Hugging Face CLI to [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache: -
-Commands - -```console -# The `delete-cache` command requires extra dependencies to work with the TUI. -# Please run `pip install huggingface_hub[cli]` to install them. - -# Launch the interactive TUI to select models to delete -$ hf delete-cache -? Select revisions to delete: 1 revisions selected counting for 438.9M. - ○ None of the following (if selected, nothing will be deleted). -Model BAAI/bge-base-en-v1.5 (438.9M, used 1 week ago) -❯ ◉ a5beb1e3: main # modified 1 week ago - -Model BAAI/bge-large-en-v1.5 (1.3G, used 1 week ago) - ○ d4aa6901: main # modified 1 week ago - -Model BAAI/bge-reranker-base (1.1G, used 4 weeks ago) - ○ 2cfc18c9: main # modified 4 weeks ago - -Press to select, to validate and to quit without modification. - -# Need to confirm after selected -? Select revisions to delete: 1 revision(s) selected. -? 1 revisions selected counting for 438.9M. Confirm deletion ? Yes -Start deletion. -Done. Deleted 1 repo(s) and 0 revision(s) for a total of 438.9M. +```bash +# delete all the cached objects +hf cache rm $(hf cache list -q) ``` -
- #### Using a proxy Here are some tips for loading/downloading models from Hugging Face using a proxy: @@ -366,15 +342,12 @@ th { | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ | | `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ | -| `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ | | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ | | `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ | -| `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ | | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ | -| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ | | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ | | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ | @@ -385,8 +358,8 @@ 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. | ✅︎ | ✅︎ | +| `DeepseekV32ForCausalLM` | DeepSeek-V3.2 | `deepseek-ai/DeepSeek-V3.2`, 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. | ✅︎ | ✅︎ | | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ | @@ -407,8 +380,8 @@ th { | `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. | ✅︎ | ✅︎ | | `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ | +| `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ | | `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ | -| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ | | `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ | | `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ | | `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ | @@ -417,13 +390,11 @@ th { | `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ | | `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ | | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | -| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ | -| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ | +| `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | | | `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. | ✅︎ | ✅︎ | | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | | @@ -443,16 +414,14 @@ th { | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ | | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ | | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ | -| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | | | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ | +| `MiniMaxM3SparseForCausalLM` | MiniMax-M3 | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ | | `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ | | `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ | | `MPTForCausalLM` | MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter | `mosaicml/mpt-7b`, `mosaicml/mpt-7b-storywriter`, `mosaicml/mpt-30b`, etc. | | ✅︎ | | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ | | `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ | -| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | -| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | | `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ | | `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | @@ -466,10 +435,8 @@ th { | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | -| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ | | `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | -| `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ | | `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | | `Qwen3ForCausalLM` | Qwen3 | `Qwen/Qwen3-8B`, etc. | ✅︎ | ✅︎ | @@ -483,26 +450,22 @@ th { | `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ | | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | | | `StableLMEpochForCausalLM` | StableLM Epoch | `stabilityai/stablelm-zephyr-3b`, etc. | | ✅︎ | -| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ | | `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ | | `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ | -| `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | | `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ | | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ | -| `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ | -| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | | -| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | | | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | | -!!! note - Grok2 requires `tokenizer.tok.json` with `tiktoken` installed. You can optionally override MoE router renormalization with `moe_router_renormalize`. - Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it! | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | +| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | +| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | +| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | !!! note Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096. @@ -544,21 +507,19 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- | | `AriaForConditionalGeneration` | Aria | T + I+ | `rhymes-ai/Aria` | | | | `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ | -| `AyaVisionForConditionalGeneration` | Aya Vision | T + I+ | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ | | `BagelForConditionalGeneration` | BAGEL | T + I+ | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | `BeeForConditionalGeneration` | Bee-8B | T + IE+ | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ | | `Blip2ForConditionalGeneration` | BLIP-2 | T + IE | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ | | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ | | `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ | | `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I+ | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ | -| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + IE+ + VE+ | `nvidia/Cosmos3-Nano` | | ✅︎ | +| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + IE+ + VE+ | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | | ✅︎ | | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I+ | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ | | `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. | | ✅︎ | @@ -579,7 +540,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ | | `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + IE+ + VE+ | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ | | `InternS2PreviewForConditionalGeneration` | Intern-S2-Preview | T + IE+ + VE+ | `internlm/Intern-S2-Preview`, etc. | ✅︎ | ✅︎ | -| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | +| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | | `KananaVForConditionalGeneration` | Kanana-V | T + I+ | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ | | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ | @@ -591,20 +552,23 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I+ | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ | | `Llama4ForConditionalGeneration` | Llama 4 | T + I+ | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ | | `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + IE+ | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ | -| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + IE+ | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ | +| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + IE+ | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ | | `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + IE+ | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ | | `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ | +| `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I+ + V+ | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | | | `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I+ + V+ | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ | | `MiDashengLMModel` | MiDashengLM | T + A+ | `mispeech/midashenglm-7b` | | ✅︎ | | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ | | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ | -| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | | +| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, `openbmb/MiniCPM-V-4_6`, etc. | ✅︎ | | +| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ | | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ | | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ | | `Molmo2ForConditionalGeneration` | Molmo2 | T + I+ / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`^, `allenai/MolmoWeb-8B`^ | ✅︎ | ✅︎ | +| `MossAudioModel` | MOSS-Audio | T + A+ | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ | +| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | -| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ | | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + IE+ + VE+ | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ | @@ -620,7 +584,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Phi4ForCausalLMV` | Phi-4-reasoning-vision | T + I+ | `microsoft/Phi-4-reasoning-vision-15B`, etc. | | ✅︎ | | `PixtralForConditionalGeneration` | Ministral 3 (Mistral format), Mistral 3 (Mistral format), Mistral Large 3 (Mistral format), Pixtral (Mistral format) | T + I+ | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, `mistralai/Mistral-Large-3-675B-Instruct-2512` `mistralai/Pixtral-12B-2409` etc. | ✅︎ | ✅︎ | | `QianfanOCRForConditionalGeneration` | QianfanOCR | T + IE+ | `baidu/Qianfan-OCR`, etc. | ✅︎ | ✅︎ | -| `QwenVLForConditionalGeneration`^ | Qwen-VL | T + IE+ | `Qwen/Qwen-VL`, `Qwen/Qwen-VL-Chat`, etc. | ✅︎ | ✅︎ | | `Qwen2AudioForConditionalGeneration` | Qwen2-Audio | T + A+ | `Qwen/Qwen2-Audio-7B-Instruct` | | ✅︎ | | `Qwen2VLForConditionalGeneration` Q | QVQ, Qwen2-VL | T + IE+ + VE+ | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ | | `Qwen2_5_VLForConditionalGeneration` Q | Qwen2.5-VL | T + IE+ + VE+ | `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen2.5-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ | @@ -637,9 +600,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Step3VLForConditionalGeneration` | Step3-VL | T + I+ | `stepfun-ai/step3` | | ✅︎ | | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I+ | `stepfun-ai/Step3-VL-10B` | | ✅︎ | | `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I+ | `stepfun-ai/Step-3.7-Flash` | | ✅︎ | -| `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ | -| `Tarsier2ForConditionalGeneration`^ | Tarsier2 | T + IE+ + VE+ | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ | | `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ | +| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I+ | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ | Some models are supported only via the [Transformers modeling backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers modeling backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it! @@ -690,9 +652,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor coordinate decoding and are not exposed by this vLLM implementation. See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes). -!!! note - To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM. - !!! note The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now. For more details, please see: @@ -711,6 +670,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `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. | ✅︎ | ✅︎ | | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ | +| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `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. | ✅︎ | ✅︎ | diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh index 4228e4954fe..d55f8c8db12 100644 --- a/docs/pre_run_check.sh +++ b/docs/pre_run_check.sh @@ -3,6 +3,26 @@ if [ "$READTHEDOCS_VERSION_TYPE" != "external" ]; then exit 0 fi +# Use a GitHub token if provided to raise the API rate limit (60 -> 5000 +# requests/hour). Set GITHUB_TOKEN in the Read the Docs environment variables. +CURL_AUTH=() +if [ -n "$GITHUB_TOKEN" ]; then + CURL_AUTH=(-H "Authorization: Bearer $GITHUB_TOKEN") +fi + +# Docs builds are now manually enabled via the 'build-docs' label. +echo "Checking for the 'build-docs' label on PR #${READTHEDOCS_VERSION_NAME}..." +LABELS=$(curl -sS "${CURL_AUTH[@]}" "https://api.github.com/repos/vllm-project/vllm/issues/${READTHEDOCS_VERSION_NAME}/labels" | python3 -c "import sys, json; print('\n'.join(l.get('name', '') for l in json.load(sys.stdin)))") +if printf '%s\n' "$LABELS" | grep -qx "build-docs"; then + echo "PR has the 'build-docs' label; continuing build." + exit 0 +else + echo "PR does not have the 'build-docs' label; cancelling build." + # See https://docs.readthedocs.com/platform/latest/guides/build/skip-build.html for info on exit code + exit 183 +fi + +# Everything below this line is effectively disabled as a temporary measure. echo "Checking for changes to docs-affecting files vs origin/main..." DOCS_PATHS=( docs/ # Actual docs content @@ -24,12 +44,6 @@ echo "Checking pre-commit/pre-run-check status..." MAX_WAIT=300 INTERVAL=60 ELAPSED=0 -# Use a GitHub token if provided to raise the API rate limit (60 -> 5000 -# requests/hour). Set GITHUB_TOKEN in the Read the Docs environment variables. -CURL_AUTH=() -if [ -n "$GITHUB_TOKEN" ]; then - CURL_AUTH=(-H "Authorization: Bearer $GITHUB_TOKEN") -fi while :; do RAW=$(curl -sS "${CURL_AUTH[@]}" -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest") HTTP_CODE=$(printf %s "$RAW" | tail -n1) diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 9fa1763108c..6d984f1a62d 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -9,12 +9,13 @@ We currently support the following OpenAI APIs: - [Completions API](./openai_compatible_server.md#completions-api) (`/v1/completions`) - Only applicable to [text generation models](../../models/generative_models.md). - *Note: `suffix` parameter is not supported.* -- [Responses API](./openai_compatible_server.md#responses-api) (`/v1/responses`) - - Only applicable to [text generation models](../../models/generative_models.md). - [Chat Completions API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions`) - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](./openai_compatible_server.md#chat-template). - *Note: `user` parameter is ignored.* - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls. +- [Chat Completions batch API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions/batch`) +- [Responses API](./openai_compatible_server.md#responses-api) (`/v1/responses`, `/v1/responses/{response_id}`, `/v1/responses/{response_id}/cancel`) + - Only applicable to [text generation models](../../models/generative_models.md). - [Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`) - Only applicable to [embedding models](../../models/pooling_models/embed.md). - [Transcriptions API](./speech_to_text.md#transcriptions-api) (`/v1/audio/transcriptions`) @@ -24,7 +25,7 @@ We currently support the following OpenAI APIs: ## Anthropic APIs -- Anthropic messages API (`/v1/messages`) +- Anthropic messages API (`/v1/messages`, `/v1/messages/count_tokens`) ## Cohere APIs @@ -35,10 +36,6 @@ We currently support the following OpenAI APIs: - Implements [Jina AI's v1 rerank API](https://jina.ai/reranker/) - compatible with [Cohere's v1 & v2 rerank APIs](https://docs.cohere.com/v2/reference/rerank) -## SageMaker APIs - -- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints) - ## Pooling APIs For further details on pooling models, please refer to [this page](../../models/pooling_models/README.md). @@ -51,7 +48,7 @@ For further details on pooling models, please refer to [this page](../../models/ - [OpenAI-compatible Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`) - Only applicable to [embedding models](../../models/pooling_models/embed.md). - [Scoring Usages](../../models/pooling_models/scoring.md) - - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`) + - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`, `/v1/score`) - [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction). - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`) @@ -68,17 +65,6 @@ For further details on speech to text, please refer to [this page](speech_to_tex - [Realtime API](./speech_to_text.md#realtime-api) (`/v1/realtime`) - Only applicable to [Automatic Speech Recognition (ASR) models](../../models/supported_models.md#realtime-transcription). -## Disaggregated APIs - -### Renderer APIs - -For further details on renderer APIs, please refer to [this page](renderer.md). - -- [Completions Render API](renderer.md) (`/v1/completions/render`) - - Render completion requests -- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`) - - Render chat completions - ## Custom APIs - [Classification API](../../models/pooling_models/classify.md#classification-api) (`/classify`) @@ -91,14 +77,79 @@ For further details on renderer APIs, please refer to [this page](renderer.md). - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`). - Computes next-token probabilities for specified `label_token_ids`. -## Utility APIs +## Instrumentator APIs + +### Basic APIs + +- `/version` - Version information +- `/load` - Server load metrics +- `/v1/models` - List available models +- `/health` - Health check + +### Metrics APIs + +For further details on metrics, please refer to [this page](../../design/metrics.md). + +- `/metrics` - Prometheus-compatible metrics HTTP endpoint + +### Offline API Documentation + +The FastAPI `/docs` endpoint requires an internet connection by default. To enable offline access in air-gapped environments, use the `--enable-offline-docs` flag: + +```bash +vllm serve NousResearch/Meta-Llama-3-8B-Instruct --enable-offline-docs +``` + +### LoRA dynamic loading + +LoRA dynamic loading & unloading is enabled in the API server. This should ONLY be used for local development! + +- `/v1/load_lora_adapter` - LoRA dynamic loading +- `/v1/unload_lora_adapter` - LoRA dynamic unloading + +### Profiling APIs + +For further details on profiling vLLM, please refer to [this page](../../contributing/profiling.md). + +- `/start_profile` - Start PyTorch profiler +- `/stop_profile` - Stop PyTorch profiler + +### SageMaker APIs + +- `/ping` - SageMaker health check +- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints) + +## Scale-Out APIs + +### Tokens IN <> Tokens OUT APIs + +- `/inference/v1/generate` - Generate completions +- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set) + +### Renderer APIs + +For further details on renderer APIs, please refer to [this page](renderer.md). + +- [Completions Render API](renderer.md) (`/v1/completions/render`) + - Render completion requests +- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`) + - Render chat completions + +### Derenderer APIs + +- `/v1/completions/derender` - Derenderer completion requests +- `/v1/chat/completions/derender` - Derenderer chat completion requests + +## Tokenize APIs - `/tokenize` - Tokenize text - `/detokenize` - Detokenize tokens -- `/health` - Health check -- `/ping` - SageMaker health check -- `/version` - Version information -- `/load` - Server load metrics +- `/tokenizer_info` - Get comprehensive tokenizer information including chat templates and configuration + +## Elastic Expert Parallelism (EEP) + +- `/scale_elastic_ep` - Trigger scaling operations +- `/is_scaling_elastic_ep` - Check if scaling is in progress ## Server in development mode @@ -119,8 +170,11 @@ For further details on Weight Transfer, please refer to [this page](../../traini - `/pause` - Pause generation (causes denial of service) - `/resume` - Resume generation - `/is_paused` - Check if generation is paused +- `/abort_requests` - Abort in-flight requests (all in-flight, or the given `request_ids`) without pausing the scheduler - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF +- `/start_weight_update` - Prepares the inference engine for a weight update. - `/update_weights` - Update model weights (can alter model behavior) +- `/finish_weight_update` - Finalizes the weight update - `/get_world_size` - Get distributed world size ### Collective RPC @@ -189,14 +243,6 @@ the detected format, which can be one of: If the result is not what you expect, you can set the `--chat-template-content-format` CLI argument to override which format to use. -## Offline API Documentation - -The FastAPI `/docs` endpoint requires an internet connection by default. To enable offline access in air-gapped environments, use the `--enable-offline-docs` flag: - -```bash -vllm serve NousResearch/Meta-Llama-3-8B-Instruct --enable-offline-docs -``` - ## Ray Serve LLM Ray Serve LLM enables scalable, production-grade serving of the vLLM engine. It integrates tightly with vLLM and extends it with features such as auto-scaling, load balancing, and back-pressure. diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index 245de012bff..e50754aa9c0 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -9,12 +9,13 @@ We currently support the following OpenAI APIs: - [Completions API](#completions-api) (`/v1/completions`) - Only applicable to [text generation models](../../models/generative_models.md). - *Note: `suffix` parameter is not supported.* -- [Responses API](#responses-api) (`/v1/responses`) - - Only applicable to [text generation models](../../models/generative_models.md). - [Chat Completions API](#chat-api) (`/v1/chat/completions`) - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](../online_serving/README.md#chat-template). - *Note: `user` parameter is ignored.* - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls. +- [Chat Completions batch API](#chat-api) (`/v1/chat/completions/batch`) +- [Responses API](#responses-api) (`/v1/responses`, `/v1/responses/{response_id}`, `/v1/responses/{response_id}/cancel`) + - Only applicable to [text generation models](../../models/generative_models.md). - [Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`) - Only applicable to [embedding models](../../models/pooling_models/embed.md). - [Transcriptions API](./speech_to_text.md#transcriptions-api) (`/v1/audio/transcriptions`) diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index d3be23fe698..e655f9c39ff 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -42,6 +42,7 @@ When using the vLLM HTTP server, the same functionality is available via: - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation +- `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) !!! note "Data Parallelism" When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. diff --git a/docs/training/layerwise.md b/docs/training/layerwise.md index d304c4a8425..5072fdb6dfd 100644 --- a/docs/training/layerwise.md +++ b/docs/training/layerwise.md @@ -28,9 +28,9 @@ For more information on implementation, see [Low Level `layerwise` API](#low-lev Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this: ```python -class Fp8OnlineLinearMethod(Fp8LinearMethod): - """Online version of Fp8LinearMethod which loads a full precision checkpoint - and quantizes weights during loading.""" +class Fp8PerTensorOnlineLinearMethod(LinearMethodBase): + """Online version of FP8 per-tensor quantization which loads a full + precision checkpoint and quantizes weights during loading.""" uses_meta_device: bool = True @@ -58,7 +58,7 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod): ### High Level Weight Transfer API -The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default. +The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle. ### Mid Level `reload_weights` API diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 4368483e8ba..7579e5fd4d0 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -17,6 +17,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba | ------- | --------- | -------- | | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference | | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU | +| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) | ## Configuration @@ -41,7 +42,7 @@ vllm serve my-model \ --weight-transfer-config '{"backend": "nccl"}' ``` -The `backend` field accepts `"nccl"` (default) or `"ipc"`. +The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`. ## API Endpoints @@ -69,7 +70,7 @@ Both backends provide static methods that the trainer calls to send weights. The EngineClass.trainer_init(init_info) # 2. Start weight update on inference side -llm.start_weight_update(is_checkpoint_format=True) +llm.start_weight_update() # 3. Send weights to inference workers EngineClass.trainer_send_weights( diff --git a/docs/training/weight_transfer/base.md b/docs/training/weight_transfer/base.md index ace228b0091..f69af8dd6e2 100644 --- a/docs/training/weight_transfer/base.md +++ b/docs/training/weight_transfer/base.md @@ -11,15 +11,23 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data ### Abstract Methods -Subclasses must implement these four methods: +Subclasses must implement these methods: | Method | Side | Description | | ------ | ---- | ----------- | | `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker | -| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally | +| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines | +| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines | +| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` | | `shutdown()` | Inference | Clean up resources | | `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process | +The base class provides two methods: + +1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`) +2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses +the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`. + ### Request Classes The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses. @@ -41,6 +49,10 @@ update_request = WeightTransferUpdateRequest( ) ``` +At the LLM/API layer, call `start_draft_weight_update()` instead of +`start_weight_update()` to target the speculative draft model; +`update_weights` / `finish_weight_update` are unchanged. + ### WeightTransferUpdateInfo The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info: @@ -81,7 +93,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo): ### 2. Implement the Engine ```python -from collections.abc import Callable, Iterator +from collections.abc import Iterator from typing import Any import torch @@ -93,18 +105,25 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]): # Set up connection to trainer using init_info.endpoint, etc. ... - def receive_weights( - self, - update_info: MyUpdateInfo, - load_weights: Callable[[list[tuple[str, torch.Tensor]]], None], - ) -> None: - # Receive each weight and call load_weights incrementally + def start_weight_update(self) -> None: + # Checkpoint-format engines: run initialize_layerwise_reload(self.model). + # In-place engines: no-op + ... + + def finish_weight_update(self) -> None: + # Checkpoint-format engines: run finalize_layerwise_reload(...). + # In-place engines: no-op + ... + + def receive_weights(self, update_info: MyUpdateInfo) -> None: + weights = [] for name, dtype_name, shape in zip( update_info.names, update_info.dtype_names, update_info.shapes ): dtype = getattr(torch, dtype_name) weight = self._fetch_weight(name, shape, dtype) - load_weights([(name, weight)]) + weights.append((name, weight)) + self.model.load_weights(weights) def shutdown(self) -> None: # Clean up resources @@ -121,9 +140,6 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]): ... ``` -!!! important - The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models. - ### 3. Register with the Factory ```python @@ -147,7 +163,7 @@ Once registered, users can select your backend via `WeightTransferConfig(backend ## WeightTransferEngineFactory -The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed. +The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed. ```python from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory @@ -155,7 +171,8 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory # Create an engine from config engine = WeightTransferEngineFactory.create_engine( config=weight_transfer_config, - parallel_config=parallel_config, + vllm_config=vllm_config, + device=device, model=model, ) ``` diff --git a/docs/training/weight_transfer/ipc.md b/docs/training/weight_transfer/ipc.md index 21fc8ad70da..f76272d2cb0 100644 --- a/docs/training/weight_transfer/ipc.md +++ b/docs/training/weight_transfer/ipc.md @@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs( llm_handle=llm_actor_handle, ) # start -ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True)) +ray.get(llm_actor_handle.start_weight_update.remote()) # send weights IPCWeightTransferEngine.trainer_send_weights( iterator=model.named_parameters(), @@ -80,7 +80,7 @@ trainer_args = IPCTrainerSendWeightsArgs( # start base_url = "http://localhost:8000" url = f"{base_url}/start_weight_update" -response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60) +response = requests.post(url, json={}, timeout=60) response.raise_for_status() # send weights IPCWeightTransferEngine.trainer_send_weights( diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index 7b531218568..481b7c5f28e 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/docs/training/weight_transfer/nccl.md @@ -11,7 +11,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b ## How It Works 1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). -2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally. +2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights. 3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py). ## Initialization @@ -93,7 +93,7 @@ remaining three steps are: from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest # 1. Start the weight update -llm.start_weight_update(is_checkpoint_format=True) +llm.start_weight_update() # 2. Receive weights (can be called multiple times for chunked transfers) llm.update_weights( @@ -116,19 +116,22 @@ must match the order in which the trainer iterates over its parameters. `start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been -transferred. The `is_checkpoint_format` flag controls whether layerwise reload -processing is applied (`True` for checkpoint-format weights, `False` for -pre-processed kernel-format weights). +transferred. The NCCL engine receives checkpoint-format weights and applies +layerwise reload processing automatically inside `start_weight_update` / +`finish_weight_update`. -Sparse NCCL patches still use `update_kind="sparse_flat"` inside -`update_info`, but they should be wrapped in -`start_weight_update(is_checkpoint_format=False)` because sparse patches apply -directly to runtime/kernel-format parameters. The current sparse MVP requires -`TP=1` and `PP=1`. +## Sparse NCCL + +Sparse, flat-index weight patches use a separate backend, +`WeightTransferConfig(backend="sparse_nccl")`, implemented by +`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group +initialization with the dense engine; patches are applied directly in place to +existing parameters (no layerwise reload). The current sparse MVP requires +`TP=1` and `PP=1`. See the example below. ## Examples - [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast -- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1` +- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1` - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model - [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane diff --git a/docs/usage/security.md b/docs/usage/security.md index 1cc91c3a8a9..ee49374e7b9 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -85,6 +85,21 @@ significantly reduce the attack surface for these types of abuse. Also, consider setting `VLLM_MEDIA_URL_ALLOW_REDIRECTS=0` to prevent HTTP redirects from being followed to bypass domain restrictions. +### 5. **Restrict Media Decode Sizes:** + +Compressed media files can expand into gigabytes of memory during decoding. vLLM +enforces decode-size limits to prevent out-of-memory denial of service: + +| Environment Variable | Default | Description | +| --- | --- | --- | +| `VLLM_MAX_IMAGE_PIXELS` | `178956970` (~179M pixels) | Maximum decoded image size in pixels. Images exceeding this are rejected before raster memory is allocated. Default matches PIL's built-in 2x decompression-bomb threshold (~680 MB for RGB). | +| `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum filesize in MB for a single audio file. | +| `VLLM_MAX_AUDIO_DECODE_DURATION_S` | `600` | Maximum decoded audio duration in seconds. Prevents compressed audio from expanding into gigabytes of float32 PCM. | + +Setting any of these to `0` disables the corresponding limit. This is **not +recommended** for deployments exposed to untrusted users, as it removes the +protection against resource-exhaustion attacks. + ## Security and Firewalls: Protecting Exposed vLLM Systems While vLLM is designed to allow unsafe network services to be isolated to @@ -176,6 +191,7 @@ The following endpoints **do not require authentication** even when `--api-key` - `/pause` - Pause generation (causes denial of service) - `/resume` - Resume generation - `/is_paused` - Check if generation is paused +- `/abort_requests` - Abort in-flight requests (causes loss of in-flight work) - `/scale_elastic_ep` - Trigger scaling operations - `/is_scaling_elastic_ep` - Check if scaling is in progress - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF @@ -311,6 +327,40 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the **Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md). +## Endpoint Plugins + +vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` entry point group (see [Endpoint Plugins](../design/endpoint_plugins.md) for how to write one). An endpoint plugin can register arbitrary FastAPI routes, including routes that reach the engine via `EngineClient.collective_rpc`, so it must be treated as part of the server's trusted code base and not as sandboxed or reviewed input. + +**Endpoint plugins are not loaded by default.** Unlike other vLLM plugin groups (`vllm.general_plugins`, `vllm.platform_plugins`, etc.), which load every discovered plugin unless `VLLM_PLUGINS` narrows the set, endpoint plugins load **none** unless `VLLM_PLUGINS` is set and explicitly names them. This mirrors the "off by default in production" posture used for development endpoints gated behind `VLLM_SERVER_DEV_MODE`. Both surfaces are only present when an operator has explicitly opted in. + +### Recommended Security Practices + +1. **Only allowlist plugins you trust.** Set `VLLM_PLUGINS` to the exact plugin names you intend to run and never wildcard or copy an allowlist between deployments without reviewing what each named plugin does. +2. **Audit routes before deploying.** A plugin's `attach_router` can add routes under any path, including ones that duplicate existing `/v1/*` paths. There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)), so a malicious or buggy plugin can **shadow a core route** and silently replace its behavior. Prefer plugins that namespace their routes under a distinct prefix (e.g. `/plugins//...`) instead of reusing `/v1/...` and review `app.routes` after startup if you need certainty about what is actually being served. +3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, and `/inference` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally. +4. **Remember the `vllm.general_plugins` pairing.** A plugin that also needs new engine side behavior ships that half separately via `vllm.general_plugins` which loads in every worker process under the default (load all unless restricted) posture. Allowlisting the endpoint plugin does not by itself restrict its paired engine side plugin. Need to review both. + +## gRPC Interface + +vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server. + +**Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network. + +### Security Implications + +An attacker who can reach the gRPC port can: + +1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials +2. **Consume GPU and compute resources** by submitting unbounded generation requests +3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM. + +### Recommendations + +- Only enable `--grpc-port` when you have a specific need for gRPC-based inference +- Ensure the gRPC port is only accessible from trusted hosts or services +- Use firewall rules to block external access to the gRPC port +- Consider deploying the gRPC interface on a dedicated internal network interface + ## Cache Directory Security vLLM assumes that its cache directories are **private and trusted**. Cache contents are loaded without cryptographic integrity verification, including formats that support arbitrary code execution. If an untrusted user or process can write to vLLM's cache directories, they may be able to crash vLLM or cause it to execute arbitrary code. diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 74d7e3eb2b0..5613d5ba4e8 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -125,10 +125,10 @@ We are working on enabling prefix caching and chunked prefill for more categorie Models using selective state-space mechanisms instead of standard transformer attention are supported. Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported. -Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`, +Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). -Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`). +Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). Please note that prefix caching is not yet supported for any of the above models. diff --git a/tests/entrypoints/serve/disagg/__init__.py b/examples/__init__.py similarity index 100% rename from tests/entrypoints/serve/disagg/__init__.py rename to examples/__init__.py diff --git a/examples/applications/chatbot/api_client.py b/examples/applications/api_server/client.py similarity index 94% rename from examples/applications/chatbot/api_client.py rename to examples/applications/api_server/client.py index 84854911bad..89207d854c9 100644 --- a/examples/applications/chatbot/api_client.py +++ b/examples/applications/api_server/client.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Example Python client for `vllm.entrypoints.api_server` +"""Example Python client for `examples/applications/api_server/server.py` Start the demo server: - python -m vllm.entrypoints.api_server --model + python examples/applications/api_server/server.py --model NOTE: The API server is used only for demonstration and simple performance benchmarks. It is not intended for production use. diff --git a/vllm/entrypoints/api_server.py b/examples/applications/api_server/server.py similarity index 99% rename from vllm/entrypoints/api_server.py rename to examples/applications/api_server/server.py index f950b52d881..adac4133210 100644 --- a/vllm/entrypoints/api_server.py +++ b/examples/applications/api_server/server.py @@ -31,7 +31,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.system_utils import set_ulimit from vllm.version import __version__ as VLLM_VERSION -logger = init_logger("vllm.entrypoints.api_server") +logger = init_logger("api_server") app = FastAPI() engine = None diff --git a/examples/applications/chatbot/gradio_webserver.py b/examples/applications/chatbot/gradio_webserver.py index f75636409c2..005bb7c68c9 100644 --- a/examples/applications/chatbot/gradio_webserver.py +++ b/examples/applications/chatbot/gradio_webserver.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Example for starting a Gradio Webserver Start vLLM API server: - python -m vllm.entrypoints.api_server \ + python examples/applications/api_server/server.py \ --model meta-llama/Llama-2-7b-chat-hf Start Webserver: diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py index 24d90eab029..cc1cc402d29 100644 --- a/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_multiturn.py @@ -35,12 +35,36 @@ Conversation isolation: the JSON body) to scope the KV cache across turns. Without it, the proxy cannot link turns and falls back to no-cache behavior. + ``conversation_id`` is a non-standard extension to the OpenAI Chat + Completions schema, consumed by this proxy and not forwarded to the + vLLM engine. Strict OpenAI-compatible frontends reject unknown + fields, so clients must opt in only when targeting this proxy. + Usage: python disagg_proxy_multiturn.py \\ --host 0.0.0.0 --port 8000 \\ --prefiller-host 10.0.0.1 --prefiller-port 8100 \\ --decoder-host 10.0.0.2 --decoder-port 8200 +Benchmarking: + Use ``benchmarks/multi_turn/benchmark_serving_multi_turn.py`` with + the ``--send-conversation-id`` flag to inject a per-conversation + ``conversation_id`` into every request so this proxy can key + cross-turn KV cache reuse. The flag is *off by default*: without + it the benchmark sends OpenAI-schema-compliant payloads and every + turn lands as a cache MISS in this proxy. + + Example: + python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\ + --model --served-model-name \\ + --url http://:8000 \\ + --input-file generate_multi_turn.json \\ + --num-clients 2 --max-active-conversations 6 \\ + --send-conversation-id + + See ``docs/features/nixl_connector_usage.md`` for the broader + bidirectional-KV-transfer setup these benchmarks exercise. + Dependencies: pip install fastapi uvicorn httpx """ @@ -373,7 +397,9 @@ async def _handle_request(api_path: str, request: Request): logger.warning( "[%s] No conversation_id provided — KV cache reuse disabled " "for this request. Add a 'conversation_id' field to enable " - "cross-turn KV sharing.", + "cross-turn KV sharing. When using " + "benchmarks/multi_turn/benchmark_serving_multi_turn.py, pass " + "--send-conversation-id (off by default).", request_id, ) diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py new file mode 100644 index 00000000000..672162f8ba3 --- /dev/null +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Push-mode disaggregated prefilling proxy demo. + +Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is +the same; the difference is in how P and D coordinate the KV transfer: + +* Pull mode: proxy forwards P's ``kv_transfer_params`` (including + ``remote_block_ids``) to D, and D pulls KV from P via NIXL READ. +* Push mode: proxy hands D **only** P's coordinates + (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``, + ``pp_size``) and the shared ``remote_request_id``. D registers its locally + allocated blocks with P over a NIXL notification; P then pushes the KV to D via + NIXL WRITE. + +Launch multiple vLLM instances configured with ``NixlPushConnector`` and +matching ``engine_id`` / ``side_channel_port``, then start this proxy: + + python3 examples/disaggregated/disaggregated_serving/\ +disagg_proxy_pushconnector_demo.py \ + --model $model_name \ + --prefill localhost:8100 \ + --decode localhost:8200 \ + --prefill-engine-id prefill-engine-001 \ + --prefill-kv-host 10.0.0.1 \ + --prefill-side-channel-port 5600 \ + --prefill-tp-size 1 \ + --prefill-pp-size 1 \ + --port 8000 +""" + +import argparse +import contextlib +import ipaddress +import itertools +import json +import logging +import os +import sys +import uuid +from abc import ABC, abstractmethod +from collections.abc import Callable + +import aiohttp +import uvicorn +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) +logger = logging.getLogger() +logging.basicConfig(level=logging.INFO) + + +class SchedulingPolicy(ABC): + @abstractmethod + def schedule(self, cycler: itertools.cycle): + raise NotImplementedError("Scheduling Proxy is not set.") + + +class RoundRobinSchedulingPolicy(SchedulingPolicy): + def schedule(self, cycler: itertools.cycle) -> str: + return next(cycler) + + +class PushProxy: + """Push-mode proxy. + + The structure mirrors the pull-mode ``Proxy`` in + ``disagg_proxy_demo.py``: an APIRouter with ``/v1/completions``, + ``/v1/chat/completions``, ``/status`` and ``/instances/add``, plus + round-robin scheduling across multiple P / D instances. + + Push-specific differences are confined to the request-handling + methods (``create_completion`` / ``create_chat_completion``): + + * D's ``kv_transfer_params`` is built from CLI-provided P + coordinates instead of being derived from P's response. + * P and D requests are issued concurrently — D registers blocks and + waits while P prefills and pushes. + """ + + def __init__( + self, + prefill_instances: list[str], + decode_instances: list[str], + model: str, + scheduling_policy: SchedulingPolicy, + prefill_engine_id: str, + prefill_kv_host: str, + prefill_side_channel_port: int, + prefill_tp_size: int, + prefill_pp_size: int, + custom_create_completion: Callable[[Request], StreamingResponse] | None = None, + custom_create_chat_completion: Callable[[Request], StreamingResponse] + | None = None, + ): + self.prefill_instances = prefill_instances + self.decode_instances = decode_instances + self.prefill_cycler = itertools.cycle(prefill_instances) + self.decode_cycler = itertools.cycle(decode_instances) + self.model = model + self.scheduling_policy = scheduling_policy + + # Push-mode metadata: D needs P's coordinates up-front. Pull mode + # learns these from P's response; push mode uses CLI args because + # D issues its registration before P responds. + self.push_metadata = { + "do_remote_decode": False, + "do_remote_prefill": True, + "remote_engine_id": prefill_engine_id, + "remote_host": prefill_kv_host, + "remote_port": prefill_side_channel_port, + "tp_size": prefill_tp_size, + "pp_size": prefill_pp_size, + } + + self.custom_create_completion = custom_create_completion + self.custom_create_chat_completion = custom_create_chat_completion + self.router = APIRouter() + self.setup_routes() + + # ── routes ──────────────────────────────────────────────────────── # + + def setup_routes(self): + self.router.post( + "/v1/completions", dependencies=[Depends(self.validate_json_request)] + )( + self.custom_create_completion + if self.custom_create_completion + else self.create_completion + ) + self.router.post( + "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)] + )( + self.custom_create_chat_completion + if self.custom_create_chat_completion + else self.create_chat_completion + ) + self.router.get("/status", response_class=JSONResponse)(self.get_status) + + async def validate_json_request(self, raw_request: Request): + content_type = raw_request.headers.get("content-type", "").lower() + if content_type != "application/json": + raise HTTPException( + status_code=415, + detail="Unsupported Media Type: Only 'application/json' is allowed", + ) + + # ── HTTP forwarding ─────────────────────────────────────────────── # + + async def forward_request(self, url, data, headers, use_chunked=True): + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + try: + async with session.post( + url=url, json=data, headers=headers + ) as response: + if 200 <= response.status < 300 or 400 <= response.status < 500: + if use_chunked: + async for chunk_bytes in response.content.iter_chunked( + 1024 + ): + yield chunk_bytes + else: + yield await response.read() + else: + error_content = await response.text() + with contextlib.suppress(json.JSONDecodeError): + error_content = json.loads(error_content) + logger.error( + "Request failed with status %s: %s", + response.status, + error_content, + ) + raise HTTPException( + status_code=response.status, + detail=f"Request failed with status {response.status}: " + f"{error_content}", + ) + except aiohttp.ClientError as e: + logger.error("ClientError occurred: %s", str(e)) + raise HTTPException( + status_code=502, + detail="Bad Gateway: Error communicating with upstream server.", + ) from e + except Exception as e: + logger.error("Unexpected error: %s", str(e)) + raise HTTPException(status_code=500, detail=str(e)) from e + + def schedule(self, cycler: itertools.cycle) -> str: + return self.scheduling_policy.schedule(cycler) + + async def get_status(self): + return { + "mode": "push", + "prefill_node_count": len(self.prefill_instances), + "decode_node_count": len(self.decode_instances), + "prefill_nodes": self.prefill_instances, + "decode_nodes": self.decode_instances, + "prefill_engine_id": self.push_metadata["remote_engine_id"], + "prefill_kv_host": self.push_metadata["remote_host"], + "prefill_side_channel_port": self.push_metadata["remote_port"], + "prefill_tp_size": self.push_metadata["tp_size"], + "prefill_pp_size": self.push_metadata["pp_size"], + } + + # ── push-mode request handling ──────────────────────────────────── # + + def _build_decode_kv_params(self, request_id: str) -> dict: + """Push-mode kv_transfer_params for D. + + ``remote_block_ids`` is intentionally omitted: D allocates its + own blocks and registers them with P; P determines the + prefill-side block IDs and ships them via the WRITE. + """ + params = self.push_metadata.copy() + params["remote_request_id"] = request_id + return params + + def _common_headers(self, request_id: str) -> dict: + h = {"X-Request-Id": request_id} + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + h["Authorization"] = f"Bearer {api_key}" + return h + + async def _push_completion(self, raw_request: Request, path: str): + """Shared body for /v1/completions and /v1/chat/completions. + + Push mode fires P and D concurrently: + * P runs a normal prefill (max_tokens=1, do_remote_decode=True). + * D runs the decode (do_remote_prefill=True, no remote_block_ids). + + D blocks waiting for P's WRITE; the response streamed back to the + client is the decode output from D. + """ + request = await raw_request.json() + request_id = str(uuid.uuid4()) + + # Prefill leg (max_tokens=1, signals P to keep KV around for D). + prefill_request = request.copy() + prefill_request["max_tokens"] = 1 + if "max_completion_tokens" in prefill_request: + prefill_request["max_completion_tokens"] = 1 + prefill_request["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None, + } + + # Decode leg (push mode: no remote_block_ids). + decode_request = request.copy() + decode_request["kv_transfer_params"] = self._build_decode_kv_params(request_id) + + prefill_instance = self.schedule(self.prefill_cycler) + decode_instance = self.schedule(self.decode_cycler) + headers = self._common_headers(request_id) + + # Fire prefill; we don't read its body but must drain the + # connection so the upstream server can free its slot. + async for _ in self.forward_request( + f"http://{prefill_instance}{path}", prefill_request, headers + ): + continue + + generator = self.forward_request( + f"http://{decode_instance}{path}", decode_request, headers + ) + return StreamingResponse(generator, media_type="application/json") + + async def create_completion(self, raw_request: Request): + try: + return await self._push_completion(raw_request, "/v1/completions") + except HTTPException: + raise + except Exception: + exc_info = sys.exc_info() + print("Error occurred in disagg push proxy server") + print(exc_info) + raise + + async def create_chat_completion(self, raw_request: Request): + try: + return await self._push_completion(raw_request, "/v1/chat/completions") + except HTTPException: + raise + except Exception: + exc_info = sys.exc_info() + error_messages = [str(e) for e in exc_info if e] + print("Error occurred in disagg push proxy server") + print(error_messages) + return StreamingResponse( + content=iter(error_messages), media_type="text/event-stream" + ) + + +class PushProxyServer: + def __init__( + self, + args: argparse.Namespace, + scheduling_policy: SchedulingPolicy | None = None, + create_completion: Callable[[Request], StreamingResponse] | None = None, + create_chat_completion: Callable[[Request], StreamingResponse] | None = None, + ): + self.validate_parsed_serve_args(args) + self.port = args.port + self.proxy_instance = PushProxy( + prefill_instances=[] if args.prefill is None else args.prefill, + decode_instances=[] if args.decode is None else args.decode, + model=args.model, + scheduling_policy=( + scheduling_policy + if scheduling_policy is not None + else RoundRobinSchedulingPolicy() + ), + prefill_engine_id=args.prefill_engine_id, + prefill_kv_host=args.prefill_kv_host, + prefill_side_channel_port=args.prefill_side_channel_port, + prefill_tp_size=args.prefill_tp_size, + prefill_pp_size=args.prefill_pp_size, + custom_create_completion=create_completion, + custom_create_chat_completion=create_chat_completion, + ) + + def validate_parsed_serve_args(self, args: argparse.Namespace): + if not args.prefill: + raise ValueError("Please specify at least one prefill node.") + if not args.decode: + raise ValueError("Please specify at least one decode node.") + if not args.prefill_engine_id: + raise ValueError( + "--prefill-engine-id is required in push mode (it must match " + "the engine_id passed to the prefill vLLM instance via " + "--kv-transfer-config)." + ) + if not args.prefill_kv_host: + raise ValueError( + "--prefill-kv-host is required in push mode (the IP / host " + "that the prefill vLLM advertises on its NIXL side channel)." + ) + self.validate_instances(args.prefill) + self.validate_instances(args.decode) + + def validate_instances(self, instances: list): + for instance in instances: + if len(instance.split(":")) != 2: + raise ValueError(f"Invalid instance format: {instance}") + host, port = instance.split(":") + try: + if host != "localhost": + ipaddress.ip_address(host) + port = int(port) + if not (0 < port < 65536): + raise ValueError(f"Invalid port number in instance: {instance}") + except Exception as e: + raise ValueError(f"Invalid instance {instance}: {str(e)}") from e + + def run_server(self): + app = FastAPI() + app.include_router(self.proxy_instance.router) + config = uvicorn.Config(app, port=self.port, loop="uvloop") + server = uvicorn.Server(config) + server.run() + + +def parse_args(): + parser = argparse.ArgumentParser("vLLM disaggregated push-mode proxy server.") + parser.add_argument("--model", "-m", type=str, required=True, help="Model name") + + parser.add_argument( + "--prefill", + "-p", + type=str, + nargs="+", + help="List of prefill node URLs (host:port)", + ) + + parser.add_argument( + "--decode", + "-d", + type=str, + nargs="+", + help="List of decode node URLs (host:port)", + ) + + parser.add_argument( + "--port", + type=int, + default=8000, + help="Server port number", + ) + + # Push-mode specific: P's coordinates that D needs in advance. + parser.add_argument( + "--prefill-engine-id", + type=str, + required=True, + help=( + "engine_id of the prefill vLLM instance (must match " + "--kv-transfer-config engine_id on the prefill server)" + ), + ) + parser.add_argument( + "--prefill-kv-host", + type=str, + required=True, + help=( + "IP / host the prefill vLLM advertises on its NIXL side " + "channel (VLLM_NIXL_SIDE_CHANNEL_HOST)" + ), + ) + parser.add_argument( + "--prefill-side-channel-port", + type=int, + default=5600, + help="NIXL side channel port on the prefill node " + "(VLLM_NIXL_SIDE_CHANNEL_PORT, default 5600)", + ) + parser.add_argument( + "--prefill-tp-size", + type=int, + default=1, + help="Tensor parallel size of the prefill vLLM instance", + ) + parser.add_argument( + "--prefill-pp-size", + type=int, + default=1, + help="Pipeline parallel size of the prefill vLLM instance", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + proxy_server = PushProxyServer(args=args) + proxy_server.run_server() diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index aceb7a9b81c..07a462711d2 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -327,6 +327,9 @@ async def handle_request(api: str, request: Request): session, decode_response = await decode_request_task stream_generator = stream_decode_response(session, decode_response, request_id) response = await make_response(stream_generator) + response.headers["Content-Type"] = decode_response.headers.get( + "Content-Type", "application/json" + ) return response except Exception as e: logger.exception("An error occurred while handling the request: %s", e) diff --git a/examples/disaggregated/lmcache/README.md b/examples/disaggregated/lmcache/README.md index 759be55d6f1..87fec826842 100644 --- a/examples/disaggregated/lmcache/README.md +++ b/examples/disaggregated/lmcache/README.md @@ -1,10 +1,38 @@ # LMCache Examples -This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing. +This folder demonstrates how to use LMCache with vLLM v1 for KV cache +offloading, disaggregated prefilling, and KV cache sharing. -## 1. Disaggregated Prefill in vLLM v1 +## Integration modes -This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node. +LMCache integrates with vLLM v1 in two ways: + +- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM + process and is configured through environment variables or a YAML config + file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node + CPU/disk offloading. +- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a + standalone server (`lmcache server`) that owns the KV cache storage; one or + more vLLM instances connect to it. This is the recommended mode for + distributed KV storage and for sharing KV cache across instances. See the + [LMCache docs](https://docs.lmcache.ai) for the full MP setup. + +## 1. CPU offload (in-process) + +- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1` + for vLLM v1. + +## 2. CPU offload (multi-process) + +- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`, + using a standalone `lmcache server`. vLLM provides a built-in shortcut for + this setup via `--kv-offloading-backend lmcache` and + `--kv-offloading-size `. + +## 3. Disaggregated Prefill in vLLM v1 + +This example demonstrates how to run LMCache with disaggregated prefill using +NIXL on a single node. ### Prerequisites @@ -46,15 +74,7 @@ The main script generates several log files: - `decoder.log` - Logs from the decode server - `proxy.log` - Logs from the proxy server -## 2. CPU Offload Examples +## 4. KV Cache Sharing -- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0 -- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1 - -## 3. KV Cache Sharing - -The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances. - -## 4. Disaggregated Prefill in vLLM v0 - -The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0. +The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV +caches between vLLM v1 instances through a centralized LMCache server. diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache.py b/examples/disaggregated/lmcache/cpu_offload_lmcache.py index 53036b3eb0f..b67a929e5d9 100644 --- a/examples/disaggregated/lmcache/cpu_offload_lmcache.py +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache.py @@ -1,20 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -This file demonstrates the example usage of cpu offloading -with LMCache in vLLM v1 or v0. - -Usage: - - Specify vLLM version - - -v v0 : Use LMCacheConnector - model = mistralai/Mistral-7B-Instruct-v0.2 - (Includes enable_chunked_prefill = True) - - -v v1 : Use LMCacheConnectorV1 (default) - model = meta-llama/Meta-Llama-3.1-8B-Instruct - (Without enable_chunked_prefill) +This file demonstrates the example usage of CPU offloading +with LMCache in vLLM v1. Note that `lmcache` is needed to run this example. Requirements: @@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to: https://docs.lmcache.ai/getting_started/installation.html """ -import argparse import contextlib import os import time @@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs def setup_environment_variables(): # LMCache-related environment variables - # Use experimental features in LMCache - os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" # LMCache is set to use 256 tokens per chunk os.environ["LMCACHE_CHUNK_SIZE"] = "256" # Enable local CPU backend in LMCache @@ -50,9 +35,9 @@ def setup_environment_variables(): @contextlib.contextmanager -def build_llm_with_lmcache(lmcache_connector: str, model: str): +def build_llm_with_lmcache(model: str): ktc = KVTransferConfig( - kv_connector=lmcache_connector, + kv_connector="LMCacheConnectorV1", kv_role="kv_both", ) # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB @@ -92,23 +77,10 @@ def print_output( print("-" * 50) -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-v", - "--version", - choices=["v0", "v1"], - default="v1", - help="Specify vLLM version (default: v1)", - ) - return parser.parse_args() - - def main(): - lmcache_connector = "LMCacheConnectorV1" model = "meta-llama/Meta-Llama-3.1-8B-Instruct" setup_environment_variables() - with build_llm_with_lmcache(lmcache_connector, model) as llm: + with build_llm_with_lmcache(model) as llm: # This example script runs two requests with a shared prefix. # Define the shared prompt and specific prompts shared_prompt = "Hello, how are you?" * 1000 diff --git a/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh b/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh new file mode 100755 index 00000000000..2372eabe1a8 --- /dev/null +++ b/examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh @@ -0,0 +1,43 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# CPU offloading with LMCache in multi-process (MP) mode. +# +# In MP mode, LMCache runs as a standalone server process (`lmcache server`) +# that owns the KV cache storage. One or more vLLM instances connect to it via +# the `LMCacheMPConnector`. This is the recommended way to run LMCache for +# distributed KV storage and for sharing KV cache across vLLM instances. +# +# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend +# lmcache` together with `--kv-offloading-size ` and vLLM wires up the +# `LMCacheMPConnector` for you (it defaults to the LMCache server at +# tcp://localhost:5555, matching the `lmcache server` default). +# +# Requires `lmcache` to be installed (`pip install lmcache`). +# Learn more: https://docs.lmcache.ai +set -euo pipefail + +MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct} + +# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by +# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache. +echo "Starting LMCache server..." +lmcache server --host localhost --port 5555 --l1-size-gb 5 & +LMCACHE_SERVER_PID=$! +trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT + +# 2. Launch vLLM and offload KV cache to the LMCache server. +# The MP connector currently requires the non-hybrid KV cache manager. +echo "Starting vLLM server with LMCache MP offloading..." +vllm serve "$MODEL" \ + --port 8000 \ + --kv-offloading-size 5 \ + --kv-offloading-backend lmcache \ + --disable-hybrid-kv-cache-manager + +# Equivalent explicit configuration (instead of the two flags above): +# --kv-transfer-config \ +# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both", +# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost", +# "lmcache.mp.port":5555}}' diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py deleted file mode 100644 index 6669eb3fb3d..00000000000 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -This file demonstrates the example usage of disaggregated prefilling -with LMCache. -We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode), -and launch an additional LMCache server. -KV cache is transferred in the following manner: -vLLM prefill node -> LMCache server -> vLLM decode node. - -Note that `pip install lmcache` is needed to run this example. -Learn more about LMCache in https://github.com/LMCache/LMCache. -""" - -import os -import subprocess -import time -from multiprocessing import Event, Process - -from lmcache.experimental.cache_engine import LMCacheEngineBuilder -from lmcache.integration.vllm.utils import ENGINE_NAME - -from vllm import LLM, SamplingParams -from vllm.config import KVTransferConfig - -# LMCache-related environment variables -# The port to start LMCache server -port = 8100 -# Use experimental features in LMCache -os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" -# LMCache is set to use 256 tokens per chunk -os.environ["LMCACHE_CHUNK_SIZE"] = "256" -# Disable local CPU backend in LMCache -os.environ["LMCACHE_LOCAL_CPU"] = "False" -# Set local CPU memory buffer limit to 5.0 GB -os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0" -# Set the remote URL for LMCache server -os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}" -# Set the serializer/deserializer between vllm and LMCache server -# `naive` indicates using raw bytes of the tensor without any compression -os.environ["LMCACHE_REMOTE_SERDE"] = "naive" - -prompts = [ - "Hello, how are you?" * 1000, -] - - -def run_prefill(prefill_done, prompts): - # We use GPU 0 for prefill node. - os.environ["CUDA_VISIBLE_DEVICES"] = "0" - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1) - - ktc = KVTransferConfig( - kv_connector="LMCacheConnector", - kv_role="kv_producer", - kv_rank=0, - kv_parallel_size=2, - ) - # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB - # memory. Reduce the value if your GPU has less memory. - llm = LLM( - model="mistralai/Mistral-7B-Instruct-v0.2", - kv_transfer_config=ktc, - max_model_len=8000, - gpu_memory_utilization=0.8, - enforce_eager=True, - ) - - # llm.generate(prompts, sampling_params) - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - generated_text = output.outputs[0].text - print(f"Generated text: {generated_text!r}") - print("Prefill node is finished.") - prefill_done.set() - - # Clean up lmcache backend - LMCacheEngineBuilder.destroy(ENGINE_NAME) - - -def run_decode(prefill_done, prompts, timeout=1): - # We use GPU 1 for decode node. - os.environ["CUDA_VISIBLE_DEVICES"] = "1" - - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) - - ktc = KVTransferConfig( - kv_connector="LMCacheConnector", - kv_role="kv_consumer", - kv_rank=1, - kv_parallel_size=2, - ) - # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB - # of memory. Reduce the value if your GPU has less memory. - llm = LLM( - model="mistralai/Mistral-7B-Instruct-v0.2", - kv_transfer_config=ktc, - max_model_len=8000, - gpu_memory_utilization=0.8, - enforce_eager=True, - ) - - print("Waiting for prefill node to finish...") - prefill_done.wait() - time.sleep(timeout) - - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - generated_text = output.outputs[0].text - print(f"Generated text: {generated_text!r}") - - # Clean up lmcache backend - LMCacheEngineBuilder.destroy(ENGINE_NAME) - - -def run_lmcache_server(port): - server_proc = subprocess.Popen( - ["python", "-m", "lmcache.experimental.server", "localhost", str(port)] - ) - return server_proc - - -def main(): - prefill_done = Event() - prefill_process = Process(target=run_prefill, args=(prefill_done, prompts)) - decode_process = Process(target=run_decode, args=(prefill_done, prompts)) - lmcache_server_process = run_lmcache_server(port) - - # Start prefill node - prefill_process.start() - - # Start decode node - decode_process.start() - - # Clean up the processes - decode_process.join() - prefill_process.terminate() - lmcache_server_process.terminate() - lmcache_server_process.wait() - - -if __name__ == "__main__": - main() diff --git a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh index 363c35028aa..61e578460c4 100644 --- a/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh +++ b/examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh @@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then UCX_TLS=cuda_ipc,cuda_copy,tcp \ LMCACHE_CONFIG_FILE=$prefill_config_file \ - LMCACHE_USE_EXPERIMENTAL=True \ VLLM_ENABLE_V1_MULTIPROCESSING=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ CUDA_VISIBLE_DEVICES=0 \ @@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then UCX_TLS=cuda_ipc,cuda_copy,tcp \ LMCACHE_CONFIG_FILE=$decode_config_file \ - LMCACHE_USE_EXPERIMENTAL=True \ VLLM_ENABLE_V1_MULTIPROCESSING=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ CUDA_VISIBLE_DEVICES=1 \ diff --git a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py index 46e2d903d4b..489ff132122 100644 --- a/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py +++ b/examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py @@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig # LMCache-related environment variables # The port to start LMCache server port = 8100 -# Use experimental features in LMCache -os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True" # LMCache is set to use 256 tokens per chunk os.environ["LMCACHE_CHUNK_SIZE"] = "256" # Disable local CPU backend in LMCache diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index b8561c73980..cfe131f000d 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -99,7 +99,7 @@ def main(): replay.send((last_seq + 1).to_bytes(8, "big")) while poller.poll(timeout=200): - seq_bytes, replay_payload = replay.recv_multipart() + _, seq_bytes, replay_payload = replay.recv_multipart() if not replay_payload: # End of replay marker is sent as an empty frame # for the payload diff --git a/examples/features/prompt_embed/prompt_embed_offline.py b/examples/features/prompt_embed/prompt_embed_offline.py index 29853bce967..9e90aa46b5b 100644 --- a/examples/features/prompt_embed/prompt_embed_offline.py +++ b/examples/features/prompt_embed/prompt_embed_offline.py @@ -19,7 +19,7 @@ Run: """ import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer +from transformers import AutoModelForCausalLM, AutoTokenizer, PythonBackend from vllm import LLM @@ -34,7 +34,7 @@ def init_tokenizer_and_llm(model_name: str): def get_prompt_embeds( chat: list[dict[str, str]], - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, embedding_layer: torch.nn.Module, ): token_ids = tokenizer.apply_chat_template( @@ -45,7 +45,7 @@ def get_prompt_embeds( def single_prompt_inference( - llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module + llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module ): chat = [{"role": "user", "content": "Please tell me about the capital of France."}] prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer) @@ -64,7 +64,7 @@ def single_prompt_inference( def batch_prompt_inference( - llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module + llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module ): chats = [ [{"role": "user", "content": "Please tell me about the capital of France."}], diff --git a/examples/features/speculative_decoding/extract_hidden_states_offline.py b/examples/features/speculative_decoding/extract_hidden_states_offline.py index f8909566f40..5db315a043b 100644 --- a/examples/features/speculative_decoding/extract_hidden_states_offline.py +++ b/examples/features/speculative_decoding/extract_hidden_states_offline.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import tempfile from vllm import LLM, SamplingParams @@ -18,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import ( with tempfile.TemporaryDirectory() as tmpdirname: llm = LLM( model="Qwen/Qwen3-8B", # Your target model - enable_chunked_prefill=False, # required speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1, @@ -38,13 +38,30 @@ with tempfile.TemporaryDirectory() as tmpdirname: kv_role="kv_producer", kv_connector_extra_config={ "shared_storage_path": tmpdirname, + "allow_custom_save_path": True, }, ), ) prompts = ["Generate a sentence with hidden states", "Write a python function"] - sampling_params = SamplingParams(max_tokens=1) - outputs = llm.generate(prompts, sampling_params) + + # One request uses defaults, the other uses a custom save path and + # includes output token hidden states via per-request kv_transfer_params. + sampling_params_list = [ + SamplingParams(max_tokens=1), + SamplingParams( + max_tokens=10, + extra_args={ + "kv_transfer_params": { + "hidden_states_path": os.path.join( + tmpdirname, "custom_output.safetensors" + ), + "include_output_tokens": True, + } + }, + ), + ] + outputs = llm.generate(prompts, sampling_params_list) for output in outputs: print("\nPrompt:", output.prompt) @@ -52,16 +69,16 @@ with tempfile.TemporaryDirectory() as tmpdirname: hidden_states_path = output.kv_transfer_params.get("hidden_states_path") assert hidden_states_path is not None - print("Prompt hidden states path:", hidden_states_path) + print("Hidden states path:", hidden_states_path) obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) token_ids = obj["token_ids"] hidden_states = obj["hidden_states"] - print("Extracted token ids:", token_ids) # Matches prompt token ids + print("Extracted token ids:", token_ids) print( "Extracted hidden states shape:", hidden_states.shape - ) # [prompt_len, num_extracted_layers, hidden_size] + ) # [num_tokens, num_extracted_layers, hidden_size] print("Extracted hidden states:", hidden_states) example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) diff --git a/examples/features/speculative_decoding/spec_decode_offline.py b/examples/features/speculative_decoding/spec_decode_offline.py index e60226ba67e..593fdb2ad00 100644 --- a/examples/features/speculative_decoding/spec_decode_offline.py +++ b/examples/features/speculative_decoding/spec_decode_offline.py @@ -72,6 +72,7 @@ def parse_args(): parser.add_argument("--max-num-seqs", type=int, default=None) parser.add_argument("--parallel-drafting", action="store_true") parser.add_argument("--allowed-local-media-path", type=str, default="") + parser.add_argument("--use-heterogeneous-vocab", action="store_true") return parser.parse_args() @@ -135,6 +136,7 @@ def main(args): "enforce_eager": args.enforce_eager, "max_model_len": args.max_model_len, "parallel_drafting": args.parallel_drafting, + "use_heterogeneous_vocab": args.use_heterogeneous_vocab, } elif args.method == "mtp": speculative_config = { diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py index c480f1b4145..fc20e8fed18 100644 --- a/examples/generate/multimodal/audio_language_offline.py +++ b/examples/generate/multimodal/audio_language_offline.py @@ -91,44 +91,6 @@ def run_cohere_asr(question: str, audio_count: int) -> ModelRequestData: ) -# MusicFlamingo -def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData: - model_name = "nvidia/music-flamingo-2601-hf" - engine_args = EngineArgs( - model=model_name, - max_model_len=4096, - max_num_seqs=2, - limit_mm_per_prompt={"audio": audio_count}, - enforce_eager=True, - ) - - # MusicFlamingo prompt placeholders use ; vLLM's MusicFlamingo - # multimodal processor expands each one into <|sound_bos|> + audio tokens + - # <|sound_eos|> based on extracted audio feature lengths. - audio_placeholder = "" * audio_count - system_prompt = ( - "You are Music Flamingo, a multimodal assistant for language and music. " - "On each turn you receive an audio clip which contains music and optional " - "text, you will receive at least one or both; use your world knowledge and " - "reasoning to help the user with any task. Interpret the entirety of the " - "content any input music--regardlenss of whether the user calls it audio, " - "music, or sound." - ) - - prompt = ( - "<|im_start|>system\n" - f"{system_prompt}<|im_end|>\n" - "<|im_start|>user\n" - f"{audio_placeholder}{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - ) - - # Gemma3N def run_gemma3n(question: str, audio_count: int) -> ModelRequestData: model_name = "google/gemma-3n-E2B-it" @@ -463,16 +425,15 @@ def run_ultravox(question: str, audio_count: int) -> ModelRequestData: # Voxtral # Make sure to install mistral-common[audio]. def run_voxtral(question: str, audio_count: int) -> ModelRequestData: - from mistral_common.audio import Audio from mistral_common.protocol.instruct.chunk import ( AudioChunk, - RawAudio, TextChunk, ) from mistral_common.protocol.instruct.messages import ( UserMessage, ) from mistral_common.protocol.instruct.request import ChatCompletionRequest + from mistral_common.tokens.tokenizers.audio import Audio from mistral_common.tokens.tokenizers.mistral import MistralTokenizer model_name = "mistralai/Voxtral-Mini-3B-2507" @@ -495,9 +456,7 @@ def run_voxtral(question: str, audio_count: int) -> ModelRequestData: Audio.from_file(str(audio_assets[i].get_local_path()), strict=False) for i in range(audio_count) ] - audio_chunks = [ - AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios - ] + audio_chunks = [AudioChunk.from_audio(audio) for audio in audios] messages = [UserMessage(content=[*audio_chunks, text_chunk])] @@ -568,7 +527,6 @@ model_example_map = { "kimi_audio": run_kimi_audio, "midashenglm": run_midashenglm, "minicpmo": run_minicpmo, - "musicflamingo": run_musicflamingo, "phi4_mm": run_phi4mm, "qwen2_audio": run_qwen2_audio, "qwen2_5_omni": run_qwen2_5_omni, diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py index 1b68a23b3bd..c3541427742 100644 --- a/examples/generate/multimodal/vision_language_multi_image_offline.py +++ b/examples/generate/multimodal/vision_language_multi_image_offline.py @@ -74,39 +74,6 @@ def load_aria(question: str, image_urls: list[str]) -> ModelRequestData: ) -def load_aya_vision(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "CohereLabs/aya-vision-8b" - - engine_args = EngineArgs( - model=model_name, - max_num_seqs=2, - limit_mm_per_prompt={"image": len(image_urls)}, - ) - - placeholders = [{"type": "image", "image": url} for url in image_urls] - messages = [ - { - "role": "user", - "content": [ - *placeholders, - {"type": "text", "text": question}, - ], - } - ] - - processor = AutoProcessor.from_pretrained(model_name) - - prompt = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - image_data=[fetch_image(url) for url in image_urls], - ) - - def load_bee(question: str, image_urls: list[str]) -> ModelRequestData: model_name = "Open-Bee/Bee-8B-RL" @@ -1042,49 +1009,6 @@ def load_phi4siglip(question: str, image_urls: list[str]) -> ModelRequestData: ) -def load_qwen_vl_chat(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "Qwen/Qwen-VL-Chat" - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=1024, - max_num_seqs=2, - hf_overrides={"architectures": ["QwenVLForConditionalGeneration"]}, - limit_mm_per_prompt={"image": len(image_urls)}, - ) - placeholders = "".join( - f"Picture {i}: \n" for i, _ in enumerate(image_urls, start=1) - ) - - # This model does not have a chat_template attribute on its tokenizer, - # so we need to explicitly pass it. We use ChatML since it's used in the - # generation utils of the model: - # https://huggingface.co/Qwen/Qwen-VL-Chat/blob/main/qwen_generation_utils.py#L265 - tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) - - # Copied from: https://huggingface.co/docs/transformers/main/en/chat_templating - chat_template = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" # noqa: E501 - - messages = [{"role": "user", "content": f"{placeholders}\n{question}"}] - prompt = tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - chat_template=chat_template, - ) - - stop_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"] - stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens] - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - stop_token_ids=stop_token_ids, - image_data=[fetch_image(url) for url in image_urls], - chat_template=chat_template, - ) - - def load_qwen2_vl(question: str, image_urls: list[str]) -> ModelRequestData: try: from qwen_vl_utils import smart_resize @@ -1318,55 +1242,6 @@ def load_step_vl(question: str, image_urls: list[str]) -> ModelRequestData: ) -def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "omni-research/Tarsier-7b" - - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=4096, - limit_mm_per_prompt={"image": len(image_urls)}, - ) - - prompt = f"USER: {'' * len(image_urls)}\n{question}\n ASSISTANT:" - image_data = [fetch_image(url) for url in image_urls] - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - image_data=image_data, - ) - - -def load_tarsier2(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "omni-research/Tarsier2-Recap-7b" - - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=32768, - limit_mm_per_prompt={"image": len(image_urls)}, - hf_overrides={ - "architectures": ["Tarsier2ForConditionalGeneration"], - "model_type": "tarsier2", - }, - ) - - prompt = ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - f"<|im_start|>user\n<|vision_start|>{'<|image_pad|>' * len(image_urls)}" - f"<|vision_end|>{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - image_data = [fetch_image(url) for url in image_urls] - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - image_data=image_data, - ) - - # GLM-4.1V def load_glm4_1v(question: str, image_urls: list[str]) -> ModelRequestData: model_name = "zai-org/GLM-4.1V-9B-Thinking" @@ -1512,7 +1387,6 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData: model_example_map = { "aria": load_aria, - "aya_vision": load_aya_vision, "bee": load_bee, "command_a_vision": load_command_a_vision, "deepseek_vl_v2": load_deepseek_vl2, @@ -1544,15 +1418,12 @@ model_example_map = { "phi4_mm": load_phi4mm, "phi4_siglip": load_phi4siglip, "pixtral_hf": load_pixtral_hf, - "qwen_vl_chat": load_qwen_vl_chat, "qwen2_vl": load_qwen2_vl, "qwen2_5_vl": load_qwen2_5_vl, "rvl": load_r_vl, "smolvlm": load_smolvlm, "step3": load_step3, "stepvl": load_step_vl, - "tarsier": load_tarsier, - "tarsier2": load_tarsier2, "glm4_1v": load_glm4_1v, "glm4_5v": load_glm4_5v, "glm4_5v_fp8": load_glm4_5v_fp8, diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 4b49d415c1b..bddf6388ae6 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -68,28 +68,6 @@ def run_aria(questions: list[str], modality: str) -> ModelRequestData: ) -# Aya Vision -def run_aya_vision(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - model_name = "CohereLabs/aya-vision-8b" - - engine_args = EngineArgs( - model=model_name, - max_model_len=2048, - max_num_seqs=2, - mm_processor_kwargs={"crop_to_patches": True}, - limit_mm_per_prompt={modality: 1}, - ) - prompts = [ - f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{question}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>" - for question in questions - ] - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Bee-8B def run_bee(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -470,24 +448,6 @@ def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData: ) -# Fuyu -def run_fuyu(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - - prompts = [f"{question}\n" for question in questions] - engine_args = EngineArgs( - model="adept/fuyu-8b", - max_model_len=2048, - max_num_seqs=2, - limit_mm_per_prompt={modality: 1}, - ) - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Gemma 3 def run_gemma3(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -1377,28 +1337,6 @@ def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData ) -# Mantis -def run_mantis(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - - llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501 - prompts = [llama3_template.format(f"{question}\n") for question in questions] - - engine_args = EngineArgs( - model="TIGER-Lab/Mantis-8B-siglip-llama3", - max_model_len=4096, - hf_overrides={"architectures": ["MantisForConditionalGeneration"]}, - limit_mm_per_prompt={modality: 1}, - ) - stop_token_ids = [128009] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - stop_token_ids=stop_token_ids, - ) - - # MiniCPM-V def run_minicpmv_base(questions: list[str], modality: str, model_name): assert modality in ["image", "video", "image+video"] @@ -1481,39 +1419,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData: return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6") -def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - - model_name = "MiniMaxAI/MiniMax-VL-01" - - engine_args = EngineArgs( - model=model_name, - max_num_seqs=2, - limit_mm_per_prompt={modality: 1}, - trust_remote_code=True, - tensor_parallel_size=8, - ) - - tokenizer = AutoTokenizer.from_pretrained(model_name) - messages = [ - [ - { - "role": "user", - "content": [{"type": "image"}, {"type": "text", "text": question}], - } - ] - for question in questions - ] - prompts = tokenizer.apply_chat_template( - messages, add_generation_prompt=True, tokenize=False - ) - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Mistral-3 HF-format def run_mistral3(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -1938,27 +1843,6 @@ def run_pixtral_hf(questions: list[str], modality: str) -> ModelRequestData: ) -# Qwen-VL -def run_qwen_vl(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - - engine_args = EngineArgs( - model="Qwen/Qwen-VL", - trust_remote_code=True, - max_model_len=1024, - max_num_seqs=2, - hf_overrides={"architectures": ["QwenVLForConditionalGeneration"]}, - limit_mm_per_prompt={modality: 1}, - ) - - prompts = [f"{question}Picture 1: \n" for question in questions] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Qwen2-VL def run_qwen2_vl(questions: list[str], modality: str) -> ModelRequestData: model_name = "Qwen/Qwen2-VL-7B-Instruct" @@ -2401,68 +2285,8 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData: ) -# omni-research/Tarsier-7b -def run_tarsier(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - model_name = "omni-research/Tarsier-7b" - - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=4096, - limit_mm_per_prompt={modality: 1}, - ) - prompts = [(f"USER: \n{question} ASSISTANT:") for question in questions] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - -def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData: - model_name = "omni-research/Tarsier2-Recap-7b" - - mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} - engine_args = EngineArgs( - model=model_name, - max_model_len=4096, - hf_overrides={ - "architectures": ["Tarsier2ForConditionalGeneration"], - "model_type": "tarsier2", - }, - limit_mm_per_prompt=mm_limit, - ) - - image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" - video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>" - - if modality == "image": - placeholder = image_placeholder - elif modality == "video": - placeholder = video_placeholder - elif modality == "image+video": - placeholder = image_placeholder + video_placeholder - - prompts = [ - ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - f"<|im_start|>user\n{placeholder}" - f"{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - for question in questions - ] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - model_example_map = { "aria": run_aria, - "aya_vision": run_aya_vision, "bagel": run_bagel, "cheers": run_cheers, "bee": run_bee, @@ -2476,7 +2300,6 @@ model_example_map = { "eagle2_5": run_eagle2_5, "ernie45_vl": run_ernie45_vl, "exaone4_5": run_exaone4_5, - "fuyu": run_fuyu, "gemma3": run_gemma3, "gemma3n": run_gemma3n, "glm4v": run_glm4v, @@ -2503,10 +2326,8 @@ model_example_map = { "llava-next": run_llava_next, "llava-next-video": run_llava_next_video, "llava-onevision": run_llava_onevision, - "mantis": run_mantis, "minicpmo": run_minicpmo, "minicpmv": run_minicpmv, - "minimax_vl_01": run_minimax_vl_01, "mistral3": run_mistral3, "molmo": run_molmo, "molmo2": run_molmo2, @@ -2522,7 +2343,6 @@ model_example_map = { "phi4_mm": run_phi4mm, "phi4_siglip": run_phi4siglip, "pixtral_hf": run_pixtral_hf, - "qwen_vl": run_qwen_vl, "qwen2_vl": run_qwen2_vl, "qwen2_5_vl": run_qwen2_5_vl, "qwen2_5_omni": run_qwen2_5_omni, @@ -2535,8 +2355,6 @@ model_example_map = { "smolvlm": run_smolvlm, "step3": run_step3, "stepvl": run_step_vl, - "tarsier": run_tarsier, - "tarsier2": run_tarsier2, } @@ -2554,15 +2372,18 @@ MODELS_NEED_VIDEO_METADATA = [ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ - "internvl_chat", + "llama4", + "qwen2_vl", "qwen2_5_vl", "qwen3_vl", "qwen3_vl_moe", - "qwen2_vl", + "kimi_vl", "qwen3_5", "qwen3_5_moe", + "internvl_chat", "stepvl", "glm4_1v", + "deepseek_ocr", ] diff --git a/examples/pooling/score/colqwen3_5_rerank_online.py b/examples/pooling/score/colqwen3_5_rerank_online.py index c64bcfc81fc..00746634d5d 100644 --- a/examples/pooling/score/colqwen3_5_rerank_online.py +++ b/examples/pooling/score/colqwen3_5_rerank_online.py @@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5. It produces per-token embeddings and uses MaxSim scoring for retrieval and reranking. Supports both text and image inputs. +Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3` +or `vultr/VultronRetrieverPrime-Qwen3.5-8B`. + Start the server with: - vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096 + vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \ + --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}' Then run this script: python colqwen3_5_rerank_online.py + +Parity note (matching the native colpali ColQwen3_5Processor pipeline): + - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792, + i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008 + (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs + as above; the default budget gives fewer visual tokens and lower retrieval ndcg. + - When you build prompts yourself (token_embed), reproduce the processor exactly: + image (document): wrap in the instruction template + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe the image.<|im_end|><|endoftext|>" + query: append the augmentation suffix + "<|endoftext|>" * 10 + Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline. """ import requests diff --git a/examples/pooling/token_classify/forced_alignment_online.py b/examples/pooling/token_classify/forced_alignment_online.py new file mode 100644 index 00000000000..01cb618e28d --- /dev/null +++ b/examples/pooling/token_classify/forced_alignment_online.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from Qwen3-ForcedAligner inference: +# https://github.com/QwenLM/Qwen3-ASR + +""" +Online forced alignment example using Qwen3-ForcedAligner-0.6B. + +Forced alignment takes audio and reference text as input and produces +word-level timestamps. The model predicts a time bin at each +token position; multiplying by ``timestamp_segment_time`` gives milliseconds. + +Start the server with: + + vllm serve Qwen/Qwen3-ForcedAligner-0.6B \\ + --runner pooling \\ + --enforce-eager \\ + --trust-request-chat-template \\ + --hf-overrides \\ + '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}' + +Then run: + + python forced_alignment_online.py +""" + +import argparse +import json +import mimetypes +import wave +from io import BytesIO +from pathlib import Path +from typing import Any + +import numpy as np +import pybase64 as base64 +import requests +import torch +from huggingface_hub import hf_hub_download + +RAW_CONTENT_CHAT_TEMPLATE = "{{ messages[0]['content'] }}" + + +def build_prompt(words: list[str]) -> str: + """Build the forced alignment prompt from a word list. + + Format: <|audio_start|><|audio_pad|><|audio_end|> + word1word2... + """ + body = "".join(words) + "" + return f"<|audio_start|><|audio_pad|><|audio_end|>{body}" + + +def encode_audio_data_uri(audio_path: Path) -> str: + mime_type = mimetypes.guess_type(audio_path)[0] or "audio/wav" + audio_base64 = base64.b64encode(audio_path.read_bytes()).decode("utf-8") + return f"data:{mime_type};base64,{audio_base64}" + + +def encode_silent_wav_data_uri(sample_rate: int = 16000, duration_s: int = 5) -> str: + audio = np.zeros(sample_rate * duration_s, dtype=np.int16) + + with BytesIO() as audio_buffer: + with wave.open(audio_buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(np.dtype(np.int16).itemsize) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio.tobytes()) + + audio_base64 = base64.b64encode(audio_buffer.getvalue()).decode("utf-8") + + return f"data:audio/wav;base64,{audio_base64}" + + +def build_payload(model: str, prompt: str, audio_uri: str) -> dict[str, Any]: + return { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "audio_url", "audio_url": {"url": audio_uri}}, + ], + } + ], + "task": "token_classify", + "chat_template": RAW_CONTENT_CHAT_TEMPLATE, + } + + +def post_http_request(payload: dict[str, Any], api_url: str) -> requests.Response: + headers = {"User-Agent": "Test Client"} + return requests.post(api_url, headers=headers, json=payload) + + +def parse_response(response: requests.Response) -> dict[str, Any]: + try: + result = response.json() + except ValueError as exc: + raise RuntimeError( + f"Server returned non-JSON response: {response.text}" + ) from exc + + if response.status_code != 200 or "data" not in result: + raise RuntimeError(f"Server error ({response.status_code}): {result}") + + return result + + +def load_timestamp_config(model: str) -> tuple[int, float]: + model_path = Path(model) + config_path = ( + model_path / "config.json" + if model_path.exists() + else Path(hf_hub_download(repo_id=model, filename="config.json")) + ) + + with config_path.open() as f: + config = json.load(f) + + return config["timestamp_token_id"], config["timestamp_segment_time"] + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="localhost") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--model", + type=str, + default="Qwen/Qwen3-ForcedAligner-0.6B", + ) + parser.add_argument( + "--audio-path", + type=Path, + default=None, + help="Optional audio file. Defaults to a 5-second silent WAV.", + ) + parser.add_argument( + "--words", + nargs="+", + default=["Hello", "world"], + help="Reference words to align against the audio.", + ) + return parser.parse_args() + + +def main(args): + from transformers import AutoTokenizer + + api_url = f"http://{args.host}:{args.port}/pooling" + prompt = build_prompt(args.words) + audio_uri = ( + encode_audio_data_uri(args.audio_path) + if args.audio_path + else encode_silent_wav_data_uri() + ) + payload = build_payload(args.model, prompt, audio_uri) + + pooling_response = post_http_request(payload=payload, api_url=api_url) + result = parse_response(pooling_response) + + tokenizer = AutoTokenizer.from_pretrained(args.model) + timestamp_token_id, timestamp_segment_time = load_timestamp_config(args.model) + + output = result["data"][0] + logits = torch.tensor(output["data"]) + predictions = logits.argmax(dim=-1) + token_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] + audio_pad_token_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>") + + usage = result.get("usage") or {} + prompt_tokens = usage.get("prompt_tokens") + if prompt_tokens is not None and prompt_tokens != len(predictions): + raise RuntimeError( + "The response length does not match the reported prompt token count." + ) + + try: + audio_pad_index = token_ids.index(audio_pad_token_id) + except ValueError as exc: + raise RuntimeError("The prompt does not contain the audio pad token.") from exc + + audio_token_shift = len(predictions) - len(token_ids) + if audio_token_shift < 0: + raise RuntimeError( + "The response is shorter than the locally tokenized prompt. " + "Check that the server was started with --trust-request-chat-template." + ) + + ts_predictions = [] + for i, token_id in enumerate(token_ids): + if token_id != timestamp_token_id: + continue + + prediction_index = i + audio_token_shift if i > audio_pad_index else i + ts_predictions.append( + predictions[prediction_index].item() * timestamp_segment_time + ) + + if len(ts_predictions) < len(args.words) * 2: + raise RuntimeError("The model did not return enough timestamp predictions.") + + for i, word in enumerate(args.words): + start_ms = ts_predictions[i * 2] + end_ms = ts_predictions[i * 2 + 1] + print(f"{word:15s} {start_ms / 1000:.3f}s - {end_ms / 1000:.3f}s") + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/pooling/token_embed/jina_reranker_v3_online.py b/examples/pooling/token_embed/jina_reranker_v3_online.py new file mode 100644 index 00000000000..8350aee2f14 --- /dev/null +++ b/examples/pooling/token_embed/jina_reranker_v3_online.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +""" +Example online usage of the Jina Reranker v3 score and rerank APIs with a task +instruction. + +Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the +server in vLLM. +""" + +import argparse +import json + +import requests + + +def post_http_request(prompt: dict, api_url: str) -> requests.Response: + headers = {"User-Agent": "Test Client"} + response = requests.post(api_url, headers=headers, json=prompt) + return response + + +def print_response(name: str, prompt: dict, response: requests.Response) -> None: + print(f"\n{name} request:") + print(json.dumps(prompt, indent=2)) + print(f"\n{name} response:") + print(json.dumps(response.json(), indent=2)) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="localhost") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3") + return parser.parse_args() + + +def main(args): + score_url = f"http://{args.host}:{args.port}/score" + rerank_url = f"http://{args.host}:{args.port}/rerank" + model_name = args.model + + query = "Which passage is about sports?" + documents = [ + "Basketball is played by two teams on a court.", + "Green tea contains antioxidants and may support metabolism.", + ] + instruction = "Rank passages about sports higher than passages about nutrition." + + score_prompt = { + "model": model_name, + "queries": query, + "documents": documents, + "instruction": instruction, + } + score_response = post_http_request(prompt=score_prompt, api_url=score_url) + print_response("Score", score_prompt, score_response) + + rerank_prompt = { + "model": model_name, + "query": query, + "documents": documents, + "instruction": instruction, + } + rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url) + print_response("Rerank", rerank_prompt, rerank_response) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/ray_serving/multi-node-serving.sh b/examples/ray_serving/multi-node-serving.sh index d2823bb8f9c..644bc820ec0 100644 --- a/examples/ray_serving/multi-node-serving.sh +++ b/examples/ray_serving/multi-node-serving.sh @@ -11,7 +11,7 @@ # Example usage: # On the head node machine, start the Ray head node process and run a vLLM server. # ./multi-node-serving.sh leader --ray_port=6379 --ray_cluster_size= [] && \ -# vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline_parallel_size 2 +# vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2 --distributed-executor-backend ray # # On each worker node, start the Ray worker node process. # ./multi-node-serving.sh worker --ray_address= --ray_port=6379 [] diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index a6adc208860..7043182ab18 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -190,12 +190,11 @@ class TrainModel: # Build platform-specific env vars for Ray -ray_env_vars = { - # Prevent Ray from setting CUDA_VISIBLE_DEVICES - "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1", -} +ray_env_vars = {} if current_platform.is_rocm(): + # Workaround for RCCL bug. See https://github.com/ROCm/rocm-systems/issues/5756 + ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1" # For ROCm, BATCH_INVARIANT vllm is not supported ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" else: @@ -307,7 +306,7 @@ gen_futures = [ ray.get(llm.pause_after_n_tokens.remote()) -ray.get(llm.start_weight_update.remote(is_checkpoint_format=True)) +ray.get(llm.start_weight_update.remote()) inference_handle = llm.update_weights.remote( WeightTransferUpdateRequest( diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py index 16c5be8dd65..0a0efcbee36 100644 --- a/examples/rl/rlhf_http_ipc.py +++ b/examples/rl/rlhf_http_ipc.py @@ -80,14 +80,10 @@ def init_weight_transfer_engine(base_url: str) -> None: response.raise_for_status() -def start_weight_update( - base_url: str, - is_checkpoint_format: bool = True, -) -> None: +def start_weight_update(base_url: str) -> None: """Start a weight update via HTTP endpoint.""" url = f"{base_url}/start_weight_update" - payload = {"is_checkpoint_format": is_checkpoint_format} - response = requests.post(url, json=payload, timeout=60) + response = requests.post(url, json={}, timeout=60) response.raise_for_status() @@ -170,7 +166,7 @@ def main(): pause_generation(BASE_URL) # Start weight update, broadcast via IPC, then finish - start_weight_update(BASE_URL, is_checkpoint_format=False) + start_weight_update(BASE_URL) print("Broadcasting weights via CUDA IPC (HTTP)...") trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL) diff --git a/examples/rl/rlhf_http_nccl.py b/examples/rl/rlhf_http_nccl.py index 01aafe43f16..b40b8de32fb 100644 --- a/examples/rl/rlhf_http_nccl.py +++ b/examples/rl/rlhf_http_nccl.py @@ -83,14 +83,10 @@ def init_weight_transfer_engine( response.raise_for_status() -def start_weight_update( - base_url: str, - is_checkpoint_format: bool = True, -) -> None: +def start_weight_update(base_url: str) -> None: """Start a weight update via HTTP endpoint.""" url = f"{base_url}/start_weight_update" - payload = {"is_checkpoint_format": is_checkpoint_format} - response = requests.post(url, json=payload, timeout=60) + response = requests.post(url, json={}, timeout=60) response.raise_for_status() @@ -223,7 +219,7 @@ def main(): shapes.append(list(p.shape)) # Start weight update - start_weight_update(BASE_URL, is_checkpoint_format=True) + start_weight_update(BASE_URL) # Start the update_weights call in a separate thread since it will block # waiting for NCCL broadcasts diff --git a/examples/rl/rlhf_ipc.py b/examples/rl/rlhf_ipc.py index afebbd240a4..cb854289879 100644 --- a/examples/rl/rlhf_ipc.py +++ b/examples/rl/rlhf_ipc.py @@ -139,7 +139,7 @@ ray.get(llm.sleep.remote(level=0)) ray.get(train_model.init_weight_transfer.remote()) # Start weight update, sync weights, then finish -ray.get(llm.start_weight_update.remote(is_checkpoint_format=True)) +ray.get(llm.start_weight_update.remote()) ray.get(train_model.broadcast_weights.remote(llm)) ray.get(llm.finish_weight_update.remote()) diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py index 0fb0a93ca82..77ac6b4cfca 100644 --- a/examples/rl/rlhf_ipc_fsdp_ep.py +++ b/examples/rl/rlhf_ipc_fsdp_ep.py @@ -277,15 +277,8 @@ class DataParallelInferenceEngine: ] ) - def start_weight_update(self, is_checkpoint_format: bool = True): - ray.get( - [ - actor.start_weight_update.remote( - is_checkpoint_format=is_checkpoint_format - ) - for actor in self.llm_actors - ] - ) + def start_weight_update(self): + ray.get([actor.start_weight_update.remote() for actor in self.llm_actors]) def finish_weight_update(self): ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors]) @@ -392,7 +385,7 @@ def main(): ray.get(inference_engine.wake_up.remote(tags=["weights"])) print("[sync] Starting weight update...") - ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True)) + ray.get(inference_engine.start_weight_update.remote()) print("[sync] Packed IPC transfer FSDP → vLLM...") ray.get( diff --git a/examples/rl/rlhf_nccl.py b/examples/rl/rlhf_nccl.py index b94d5e4db82..bebd6bc70df 100644 --- a/examples/rl/rlhf_nccl.py +++ b/examples/rl/rlhf_nccl.py @@ -29,6 +29,7 @@ causes unexpected behavior. import os import ray +import torch from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from transformers import AutoModelForCausalLM @@ -39,12 +40,24 @@ from vllm.distributed.weight_transfer.nccl_engine import ( NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine, ) +from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, get_open_port MODEL_NAME = "facebook/opt-125m" # MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128" +def get_assigned_gpu(): + """This is a temporary workaround for a runtime bug in RCCL on ROCm.""" + if not current_platform.is_rocm(): + return 0 + assigned_gpu = int(ray.get_gpu_ids()[0]) + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + os.environ.pop("HIP_VISIBLE_DEVICES", None) + torch.accelerator.set_device_idx(assigned_gpu) + return assigned_gpu + + class MyLLM(LLM): """Configure the vLLM worker for Ray placement group execution.""" @@ -58,9 +71,11 @@ class TrainModel: """Ray actor that wraps the training model on a dedicated GPU.""" def __init__(self, model_name: str): + assigned_gpu = get_assigned_gpu() + self.model = AutoModelForCausalLM.from_pretrained( model_name, - ).to("cuda:0") + ).to(f"cuda:{assigned_gpu}") self.port = get_open_port() self.master_address = get_ip() @@ -187,7 +202,7 @@ ray.get([train_handle, inference_handle]) names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote()) # Start weight update -ray.get(llm.start_weight_update.remote(is_checkpoint_format=True)) +ray.get(llm.start_weight_update.remote()) # Issue update_weights call with NCCL-specific update info # packed=True enables efficient batched tensor broadcasting diff --git a/examples/rl/rlhf_nccl_fsdp_ep.py b/examples/rl/rlhf_nccl_fsdp_ep.py index 4337e6fea5a..860db058cac 100644 --- a/examples/rl/rlhf_nccl_fsdp_ep.py +++ b/examples/rl/rlhf_nccl_fsdp_ep.py @@ -299,7 +299,7 @@ async def main(): print(f"[sync] Got metadata for {len(names)} parameters.") print("[sync] Starting weight update...") - await engine.start_weight_update(is_checkpoint_format=True) + await engine.start_weight_update() print("[sync] Broadcasting weights from FSDP → vLLM...") broadcast_handles = [ diff --git a/examples/rl/rlhf_sparse_nccl.py b/examples/rl/rlhf_sparse_nccl.py index bddd28b6485..09cf5bfbaa0 100644 --- a/examples/rl/rlhf_sparse_nccl.py +++ b/examples/rl/rlhf_sparse_nccl.py @@ -44,11 +44,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.base import SparseWeightPatch from vllm.distributed.weight_transfer.nccl_engine import ( NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine, ) +from vllm.distributed.weight_transfer.sparse_nccl_engine import ( + SparseNCCLWeightTransferEngine, + SparseWeightPatch, +) from vllm.utils.network_utils import get_ip, get_open_port MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" @@ -244,7 +247,6 @@ class TrainModel: dtype_names=[str(self.patched_param.dtype).split(".")[-1]], shapes=[list(self.patched_param.shape)], num_updates_list=[flat_indices.numel()], - update_kind="sparse_flat", ) return update_info, selected_token_ids, patch_digest, sparse_payload_bytes @@ -271,7 +273,7 @@ class TrainModel: raise RuntimeError("Sparse patch has not been prepared") start = time.perf_counter() - NCCLWeightTransferEngine.trainer_send_sparse_weights( + SparseNCCLWeightTransferEngine.trainer_send_weights( iter(self.pending_sparse_patches), NCCLTrainerSendWeightsArgs(group=self.model_update_group), ) @@ -282,6 +284,7 @@ class TrainModel: def launch_llm( scheduling_inference: PlacementGroupSchedulingStrategy, + backend: str = "nccl", ): return ray.remote( num_cpus=0, @@ -293,7 +296,7 @@ def launch_llm( tensor_parallel_size=1, distributed_executor_backend="ray", gpu_memory_utilization=0.7, - weight_transfer_config=WeightTransferConfig(backend="nccl"), + weight_transfer_config=WeightTransferConfig(backend=backend), ) @@ -332,7 +335,7 @@ def run_dense_phase( scheduling_inference: PlacementGroupSchedulingStrategy, ) -> dict[str, object]: ray.get(train_model.reset_model.remote()) - llm = launch_llm(scheduling_inference) + llm = launch_llm(scheduling_inference, backend="nccl") try: dense_before = collect_vllm_generations(llm) @@ -351,7 +354,7 @@ def run_dense_phase( ) trainer_init = train_model.init_weight_transfer_group.remote(world_size) ray.get([trainer_init, inference_init]) - ray.get(llm.start_weight_update.remote(is_checkpoint_format=True)) + ray.get(llm.start_weight_update.remote()) dense_update_info, dense_payload_bytes = ray.get( train_model.get_dense_update_info.remote() @@ -391,7 +394,7 @@ def run_sparse_phase( scheduling_inference: PlacementGroupSchedulingStrategy, ) -> dict[str, object]: ray.get(train_model.reset_model.remote()) - llm = launch_llm(scheduling_inference) + llm = launch_llm(scheduling_inference, backend="sparse_nccl") try: sparse_before = collect_vllm_generations(llm) @@ -410,7 +413,7 @@ def run_sparse_phase( ) trainer_init = train_model.init_weight_transfer_group.remote(world_size) ray.get([trainer_init, inference_init]) - ray.get(llm.start_weight_update.remote(is_checkpoint_format=False)) + ray.get(llm.start_weight_update.remote()) sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = ( ray.get(train_model.prepare_sparse_patch.remote(PROMPTS)) diff --git a/tests/entrypoints/serve/render/__init__.py b/examples/scale_out/__init__.py similarity index 100% rename from tests/entrypoints/serve/render/__init__.py rename to examples/scale_out/__init__.py diff --git a/examples/disaggregated/disaggregated_serving/example_mm_serve.py b/examples/scale_out/example_mm_serve.py similarity index 100% rename from examples/disaggregated/disaggregated_serving/example_mm_serve.py rename to examples/scale_out/example_mm_serve.py diff --git a/examples/generate/token_generation_client.py b/examples/scale_out/token_generation_client.py similarity index 100% rename from examples/generate/token_generation_client.py rename to examples/scale_out/token_generation_client.py diff --git a/examples/speech_to_text/openai/openai_transcription_client.py b/examples/speech_to_text/openai/openai_transcription_client.py index 396edba1155..f928c06d45e 100644 --- a/examples/speech_to_text/openai/openai_transcription_client.py +++ b/examples/speech_to_text/openai/openai_transcription_client.py @@ -33,15 +33,23 @@ def sync_openai( *, repetition_penalty: float = 1.3, hotwords: str = None, + prompt: str | None = None, ): """ Perform synchronous transcription using OpenAI-compatible API. + + The optional ``prompt`` is the OpenAI-API ``prompt`` field (style / + vocabulary hint). It is wired through model-by-model: Whisper uses it + as a ``<|prev|>`` continuation hint, Qwen3-ASR maps it into the + chat-template ``system`` turn. Models that do not consume it accept + it without effect. """ with open(audio_path, "rb") as f: transcription = client.audio.transcriptions.create( file=f, model=model, language="en", + prompt=prompt or "", response_format="json", temperature=0.0, # Additional sampling params not provided by OpenAI API. @@ -55,7 +63,11 @@ def sync_openai( async def stream_openai_response( - audio_path: str, client: AsyncOpenAI, model: str, hotwords: str = None + audio_path: str, + client: AsyncOpenAI, + model: str, + hotwords: str = None, + prompt: str | None = None, ): """ Perform asynchronous transcription using OpenAI-compatible API. @@ -66,6 +78,7 @@ async def stream_openai_response( file=f, model=model, language="en", + prompt=prompt or "", response_format="json", temperature=0.0, # Additional sampling params not provided by OpenAI API. @@ -146,6 +159,7 @@ def main(args): model=model, repetition_penalty=args.repetition_penalty, hotwords=args.hotwords, + prompt=args.prompt, ) # Run the asynchronous function @@ -160,6 +174,7 @@ def main(args): client, model, hotwords=args.hotwords, + prompt=args.prompt, ) ) else: @@ -193,5 +208,16 @@ if __name__ == "__main__": default=None, help="hotwords", ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help=( + "Optional `prompt` (OpenAI transcription API: style/vocabulary " + "hint). Wired model-by-model: Whisper uses it as a `<|prev|>` " + "continuation hint, Qwen3-ASR maps it into the chat-template " + "system turn." + ), + ) args = parser.parse_args() main(args) diff --git a/examples/template_baichuan.jinja b/examples/template_baichuan.jinja deleted file mode 100644 index 42a8d9270a4..00000000000 --- a/examples/template_baichuan.jinja +++ /dev/null @@ -1,13 +0,0 @@ -{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }} - -{%- for message in messages -%} - {%- if message['role'] == 'user' -%} - {{- '' + message['content'] -}} - {%- elif message['role'] == 'assistant' -%} - {{- '' + message['content'] -}} - {%- endif -%} -{%- endfor -%} - -{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%} - {{- '' -}} -{% endif %} \ No newline at end of file diff --git a/examples/tool_chat_template_gemma4.jinja b/examples/tool_chat_template_gemma4.jinja index ef765823106..6ce01e6479a 100644 --- a/examples/tool_chat_template_gemma4.jinja +++ b/examples/tool_chat_template_gemma4.jinja @@ -116,7 +116,9 @@ } {%- endmacro -%} {%- macro format_argument(argument, escape_keys=True) -%} - {%- if argument is string -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} {{- '<|"|>' + argument + '<|"|>' -}} {%- elif argument is boolean -%} {{- 'true' if argument else 'false' -}} @@ -172,18 +174,21 @@ {{- '' -}} {%- endmacro -%} -{%- set ns = namespace(prev_message_type=None) -%} +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} {%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{%- set preserve_thinking = preserve_thinking | default(false) -%} {{- bos_token -}} {#- Handle System/Tool Definitions Block -#} -{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} {{- '<|turn>system\n' -}} {#- Inject Thinking token at the very top of the FIRST system turn -#} - {%- if enable_thinking is defined and enable_thinking -%} + {%- if enable_thinking -%} {{- '<|think|>\n' -}} {%- set ns.prev_message_type = 'think' -%} {%- endif -%} - {%- if messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} {%- if messages[0]['content'] is string -%} {{- messages[0]['content'] | trim -}} {%- elif messages[0]['content'] is sequence -%} @@ -217,31 +222,24 @@ {%- if message['role'] != 'tool' -%} {%- set ns.prev_message_type = None -%} {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} - {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#} - {%- set prev_nt = namespace(role=None, found=false) -%} - {%- if loop.index0 > 0 -%} - {%- for j in range(loop.index0 - 1, -1, -1) -%} - {%- if not prev_nt.found -%} - {%- if loop_messages[j]['role'] != 'tool' -%} - {%- set prev_nt.role = loop_messages[j]['role'] -%} - {%- set prev_nt.found = true -%} - {%- endif -%} - {%- endif -%} - {%- endfor -%} - {%- endif -%} - {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%} + {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} {%- if not continue_same_model_turn -%} {{- '<|turn>' + role + '\n' }} + {%- if role == 'model' and not enable_thinking and not (message.get('reasoning') or message.get('reasoning_content')) -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} {%- endif -%} {#- Render reasoning/reasoning_content as thinking channel -#} {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} - {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%} + {%- if thinking_text and thinking_gate -%} {{- '<|channel>thought\n' + thinking_text + '\n' -}} {%- endif -%} - {%- if message['tool_calls'] -%} - {%- for tool_call in message['tool_calls'] -%} + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} {%- set function = tool_call['function'] -%} {{- '<|tool_call>call:' + function['name'] + '{' -}} {%- if function['arguments'] is mapping -%} @@ -251,8 +249,13 @@ {%- set ns_args.found_first = true -%} {{- key -}}:{{- format_argument(value, escape_keys=False) -}} {%- endfor -%} - {%- elif function['arguments'] is string -%} - {{- function['arguments'] -}} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} {%- endif -%} {{- '}' -}} {%- endfor -%} @@ -262,7 +265,7 @@ {%- set ns_tr_out = namespace(flag=false) -%} {%- if message.get('tool_responses') -%} {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} - {%- for tool_response in message['tool_responses'] -%} + {%- for tool_response in message.get('tool_responses') -%} {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} {%- set ns_tr_out.flag = true -%} {%- set ns.prev_message_type = 'tool_response' -%} @@ -277,8 +280,8 @@ {%- else -%} {%- set follow = loop_messages[k] -%} {#- Resolve tool_call_id to function name -#} - {%- set ns_tname = namespace(name=follow.get('name') | default('unknown', true)) -%} - {%- for tc in message['tool_calls'] -%} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} {%- if tc.get('id') == follow.get('tool_call_id') -%} {%- set ns_tname.name = tc['function']['name'] -%} {%- endif -%} @@ -296,9 +299,9 @@ {%- endfor -%} {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} {%- for part in tool_body -%} - {%- if part.get('type') == 'image' -%} + {%- if part.get('type') in ['image', 'image_url'] -%} {{- '<|image|>' -}} - {%- elif part.get('type') == 'audio' -%} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} {{- '<|audio|>' -}} {%- elif part.get('type') == 'video' -%} {{- '<|video|>' -}} @@ -314,29 +317,26 @@ {%- endif -%} {%- set captured_content -%} - {%- if message['content'] is string -%} + {%- if message.get('content') is string -%} {%- if role == 'model' -%} {{- strip_thinking(message['content']) -}} {%- else -%} {{- message['content'] | trim -}} {%- endif -%} - {%- elif message['content'] is sequence -%} + {%- elif message.get('content') is sequence -%} {%- for item in message['content'] -%} - {%- if item['type'] == 'text' -%} + {%- if item.get('type') == 'text' -%} {%- if role == 'model' -%} {{- strip_thinking(item['text']) -}} {%- else -%} {{- item['text'] | trim -}} {%- endif -%} - {%- elif item['type'] == 'image' -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} {{- '<|image|>' -}} - {%- set ns.prev_message_type = 'image' -%} - {%- elif item['type'] == 'audio' -%} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} {{- '<|audio|>' -}} - {%- set ns.prev_message_type = 'audio' -%} - {%- elif item['type'] == 'video' -%} + {%- elif item.get('type') == 'video' -%} {{- '<|video|>' -}} - {%- set ns.prev_message_type = 'video' -%} {%- endif -%} {%- endfor -%} {%- endif -%} @@ -345,19 +345,43 @@ {{- captured_content -}} {%- set has_content = captured_content | trim | length > 0 -%} + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and (not message.get('tool_calls') or ns_tr_out.flag) + ) -%} + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {{- '\n' -}} {%- elif not (ns_tr_out.flag and not has_content) -%} {{- '\n' -}} {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} {%- endif -%} {%- endfor -%} {%- if add_generation_prompt -%} {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} {{- '<|turn>model\n' -}} - {%- if not enable_thinking | default(false) -%} + {%- if not enable_thinking -%} {{- '<|channel>thought\n' -}} {%- endif -%} + {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%} + {{- '<|channel>thought\n' -}} {%- endif -%} -{%- endif -%} \ No newline at end of file +{%- endif -%} diff --git a/examples/tool_chat_template_minimax_m1.jinja b/examples/tool_chat_template_minimax_m1.jinja deleted file mode 100644 index 2d5bbf4de56..00000000000 --- a/examples/tool_chat_template_minimax_m1.jinja +++ /dev/null @@ -1,91 +0,0 @@ -{{ '' -}} -{%- if custom_tools is defined %} - {%- set tools = custom_tools %} -{%- endif %} -{%- if not tools is defined %} - {%- set tools = none %} -{%- endif %} - -{#- Extract system message #} -{% set ns = namespace(system_prompt='') -%} -{%- if messages[0]['role'] == 'system' %} - {%- if messages[0]['content'] is string %} - {%- set ns.system_prompt = messages[0]['content']|trim %} - {%- else %} - {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %} - {%- endif %} - {%- set messages = messages[1:] %} -{%- else %} - {%- if tools is not none %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- else %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- endif %} -{%- endif %} - -{#- System message #} -{%- if ns.system_prompt != '' %} -{{ 'system ai_setting=assistant\n' + ns.system_prompt + '\n' -}} -{%- endif %} - -{#- Tools configuration #} -{%- if tools is not none %} -{{ 'system tool_setting=tools\nYou are provided with these tools:\n\n' -}} -{%- for tool in tools %} -{{ tool | tojson ~ '\n' -}} -{%- endfor %} -{{ '\n\nIf you need to call tools, please respond with XML tags, and provide tool-name and json-object of arguments, following the format below:\n\n{"name": , "arguments": }\n...\n\n' -}} -{%- endif %} - -{#- Process messages #} -{%- for message in messages %} - {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %} - {%- if message['role'] == 'user' %} -{{ 'user name=user\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ content['text']|trim -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- elif message['role'] == 'assistant' %} -{{ 'ai name=assistant\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %} -{{ content['text']|trim -}} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} - {%- elif 'tool_calls' in message %} -{{ 'ai name=assistant\n\n' -}} -{%- for tool_call in message.tool_calls %} -{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}} -{%- endfor %} -{{ '\n' -}} - {%- elif message.role == "tool" or message.role == "ipython" %} -{{ 'tool name=tools\n' -}} -{%- if message.content is string %} -{{ 'tool result: ' + message.content + '\n\n' -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ 'tool result: ' + content['text'] + '\n\n' -}} -{%- elif content.get('name') %} -{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} -{%- endfor %} - -{%- if add_generation_prompt %} -{{ 'ai name=assistant\n' -}} -{%- endif %} \ No newline at end of file diff --git a/mkdocs.yaml b/mkdocs.yaml index 970bf963309..a32cea61806 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -114,6 +114,7 @@ plugins: features/quantization/int4.md: features/quantization/llm_compressor/int4.md features/quantization/int8.md: features/quantization/llm_compressor/int8_w8a8.md serving/openai_compatible_server.md: serving/online_serving/README.md + examples/others/lmcache.md: examples/disaggregated/lmcache.md markdown_extensions: - attr_list diff --git a/pyproject.toml b/pyproject.toml index c782cc326bc..3819ad7fc8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,8 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/tool-parser/src/gemma4.rs", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs"] + "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/text/src/output/decoded.rs", + "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] ignore-hidden = false [tool.typos.default] @@ -162,6 +163,8 @@ dout = "dout" Pn = "Pn" arange = "arange" thw = "thw" +# temporal position ids (parallels hpos/wpos in vision RoPE) +tpos = "tpos" subtile = "subtile" HSA = "HSA" setp = "setp" diff --git a/requirements/build/cpu.txt b/requirements/build/cpu.txt index 640432ddd8c..27a3ac65c98 100644 --- a/requirements/build/cpu.txt +++ b/requirements/build/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu cmake>=3.26.1 ninja packaging>=24.2 diff --git a/requirements/build/rust.txt b/requirements/build/rust.txt new file mode 100644 index 00000000000..e2874dee0ab --- /dev/null +++ b/requirements/build/rust.txt @@ -0,0 +1,4 @@ +# Dependencies for building Rust artifacts through setuptools-rust. +setuptools>=77.0.3,<81.0.0 +setuptools-rust>=1.9.0 +wheel diff --git a/requirements/common.txt b/requirements/common.txt index 8b37f3cd30c..30332ed448b 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -7,32 +7,33 @@ requests >= 2.26.0 tqdm blake3 py-cpuinfo -transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 +transformers >= 5.5.3 tokenizers >= 0.21.1 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 -fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. +fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides. +starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1 aiohttp >= 3.13.3 openai >= 2.0.0 # For Responses API with reasoning content pydantic >= 2.12.0 prometheus_client >= 0.18.0 pillow # Required for image processing -prometheus-fastapi-instrumentator >= 7.0.0 +prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0 tiktoken >= 0.6.0 # Required for DBRX tokenizer lm-format-enforcer == 0.11.3 -llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" +llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x" outlines_core == 0.2.14 # required for outlines backend disk cache diskcache == 5.6.3 lark == 1.2.2 -xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" +xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" typing_extensions >= 4.10 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317 partial-json-parser # used for parsing partial JSON outputs +jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation pyzmq >= 25.0.0 msgspec -gguf >= 0.17.0 -mistral_common[image] >= 1.11.3 +mistral_common[image] >= 1.11.5 opencv-python-headless >= 4.13.0 # required for video IO pyyaml six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12 diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 5ec338af736..c0b98d22c9b 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu # Common dependencies -r common.txt @@ -16,6 +15,9 @@ torchaudio; platform_machine != "s390x" and platform_machine != "riscv64" # required for the image processor of phi3v, this must be updated alongside torch torchvision; platform_machine != "s390x" and platform_machine != "riscv64" +# required for the torchcodec video decoding backend +torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64" and platform_machine != "ppc64le" + # Intel Extension for PyTorch, only for x86_64 CPUs intel-openmp==2024.2.1; platform_machine == "x86_64" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 89be67be8f5..116aca3c7c6 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -8,13 +8,16 @@ torch==2.11.0 torchaudio==2.11.0 # These must be updated alongside torch torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version +torchcodec >= 0.14 +PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.12 -flashinfer-cubin==0.6.12 +flashinfer-python==0.6.13 +flashinfer-cubin==0.6.13 apache-tvm-ffi==0.1.9 tilelang==0.1.9 nvidia-cudnn-frontend>=1.19.1 - +# Required for LLM_NVTX_SCOPES_FOR_PROFILING=1 +nvtx==0.2.15 # Required for faster safetensors model loading fastsafetensors >= 0.3.2 @@ -23,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.2 +tokenspeed-mla==0.1.2; platform_system == "Linux" # Humming kernels for quantization gemm -humming-kernels[cu13]==0.1.4 +humming-kernels[cu13]==0.1.10 diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index 7a5b5f25c37..ce920816db3 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,5 +2,5 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0 # until a fixed newer release is verified for runtime images. cupy-cuda13x < 14.1.0 -nixl >= 1.1.0 # Required for disaggregated prefill +nixl == 1.3.0 mooncake-transfer-engine >= 0.3.8 diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 4ca70738303..5179f6ee8d7 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -23,6 +23,7 @@ timm>=1.0.17 # To be consistent with test_quark.py amd-quark>=0.8.99 tilelang==0.1.10 - +# Required apache-tvm-ffi matching tilelang version +apache-tvm-ffi==0.1.10 # Required for faster safetensors model loading fastsafetensors >= 0.3.2 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt new file mode 100644 index 00000000000..b6918202ac7 --- /dev/null +++ b/requirements/test/cpu.txt @@ -0,0 +1,1285 @@ +# This file was autogenerated by uv via the following command: +# uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_28 --python-version 3.12 +absl-py==2.1.0 + # via rouge-score +accelerate==1.13.0 + # via peft +aiohappyeyeballs==2.6.1 + # via aiohttp +aiohttp==3.13.3 + # via + # -r requirements/test/../common.txt + # aiohttp-cors + # datasets + # fsspec + # gpt-oss + # lm-eval + # ray +aiohttp-cors==0.8.1 + # via ray +aiosignal==1.4.0 + # via aiohttp +albumentations==1.4.6 + # via -r requirements/test/cuda.in +alembic==1.16.4 + # via optuna +annotated-doc==0.0.4 + # via + # fastapi + # typer +annotated-types==0.7.0 + # via pydantic +anthropic==0.112.0 + # via -r requirements/test/../common.txt +anyio==4.14.1 + # via + # anthropic + # httpx + # mcp + # openai + # sse-starlette + # starlette + # watchfiles +apache-tvm-ffi==0.1.9 + # via xgrammar +arctic-inference==0.1.1 + # via -r requirements/test/cuda.in +argcomplete==3.5.1 + # via datamodel-code-generator +astor==0.8.1 + # via depyf +attrs==24.2.0 + # via + # aiohttp + # hypothesis + # jsonschema + # referencing +audioread==3.0.1 + # via librosa +av==16.1.0 + # via -r requirements/test/cuda.in +azure-core==1.38.2 + # via + # azure-identity + # azure-storage-blob +azure-identity==1.25.2 + # via runai-model-streamer-azure +azure-storage-blob==12.28.0 + # via runai-model-streamer-azure +backoff==2.2.1 + # via -r requirements/test/cuda.in +bitsandbytes==0.49.2 + # via -r requirements/test/cuda.in +black==24.10.0 + # via datamodel-code-generator +blake3==1.0.9 + # via -r requirements/test/../common.txt +blobfile==3.0.0 + # via -r requirements/test/cuda.in +bm25s==0.2.13 + # via mteb +boto3==1.35.57 + # via + # runai-model-streamer-s3 + # tensorizer +botocore==1.35.57 + # via + # boto3 + # s3transfer +bounded-pool-executor==0.0.3 + # via pqdm +buildkite-test-collector==0.1.9 + # via -r requirements/test/cuda.in +cachetools==5.5.2 + # via + # -r requirements/test/../common.txt + # google-auth +cbor2==6.1.2 + # via -r requirements/test/../common.txt +certifi==2024.8.30 + # via + # httpcore + # httpx + # requests + # sentry-sdk +cffi==2.0.0 + # via + # cryptography + # soundfile +chardet==5.2.0 + # via mbstrdecoder +charset-normalizer==3.4.0 + # via requests +chz==0.3.0 + # via gpt-oss +click==8.4.2 + # via + # black + # huggingface-hub + # jiwer + # nltk + # ray + # rich-toolkit + # schemathesis + # uvicorn +cloudpickle==3.1.2 + # via -r requirements/test/../common.txt +cohere-melody==0.9.0 + # via -r requirements/test/cuda.in +colorama==0.4.6 + # via + # perceptron + # sacrebleu +colorful==0.5.6 + # via ray +colorlog==6.10.1 + # via optuna +compressed-tensors==0.17.0 + # via -r requirements/test/../common.txt +contourpy==1.3.0 + # via matplotlib +coverage==7.10.6 + # via pytest-cov +cramjam==2.9.0 + # via fastparquet +cryptography==46.0.5 + # via + # azure-identity + # azure-storage-blob + # msal + # pyjwt +cupy-cuda12x==13.6.0 + # via ray +cycler==0.12.1 + # via matplotlib +datamodel-code-generator==0.26.3 + # via -r requirements/test/cuda.in +dataproperty==1.0.1 + # via + # pytablewriter + # tabledata +datasets==3.3.0 + # via + # -r requirements/test/cuda.in + # evaluate + # lm-eval + # mteb +decorator==5.1.1 + # via librosa +decord==0.6.0 + # via -r requirements/test/cuda.in +depyf==0.20.0 + # via -r requirements/test/../common.txt +detect-installer==0.1.0 + # via fastapi-cloud-cli +dill==0.3.8 + # via + # datasets + # depyf + # evaluate + # lm-eval + # multiprocess +diskcache==5.6.3 + # via -r requirements/test/../common.txt +distlib==0.3.9 + # via virtualenv +distro==1.9.0 + # via + # anthropic + # openai +dnspython==2.7.0 + # via email-validator +docker==7.1.0 + # via gpt-oss +docopt==0.6.2 + # via num2words +docstring-parser==0.18.0 + # via anthropic +einops==0.8.1 + # via + # -r requirements/test/../common.txt + # encodec + # vector-quantize-pytorch + # vocos +einx==0.3.0 + # via vector-quantize-pytorch +email-validator==2.2.0 + # via + # fastapi + # pydantic +encodec==0.1.1 + # via vocos +et-xmlfile==2.0.0 + # via openpyxl +evaluate==0.4.3 + # via lm-eval +fastapi==0.136.3 + # via + # -r requirements/test/../common.txt + # gpt-oss + # model-hosting-container-standards +fastapi-cli==0.0.27 + # via fastapi +fastapi-cloud-cli==0.21.0 + # via fastapi-cli +fastar==0.11.0 + # via + # fastapi + # fastapi-cloud-cli +fastparquet==2024.11.0 + # via genai-perf +fastrlock==0.8.2 + # via cupy-cuda12x +fastsafetensors==0.3.2 + # via -r requirements/test/cuda.in +filelock==3.16.1 + # via + # -r requirements/test/../common.txt + # blobfile + # datasets + # huggingface-hub + # ray + # torch + # virtualenv +fonttools==4.55.0 + # via matplotlib +frozendict==2.4.6 + # via einx +frozenlist==1.5.0 + # via + # aiohttp + # aiosignal +fsspec==2024.12.0 + # via + # datasets + # evaluate + # fastparquet + # huggingface-hub + # torch +ftfy==6.3.1 + # via open-clip-torch +genai-perf==0.0.16 + # via -r requirements/test/cuda.in +genson==1.3.0 + # via datamodel-code-generator +google-api-core==2.24.2 + # via + # google-cloud-core + # google-cloud-storage + # opencensus +google-auth==2.40.2 + # via + # google-api-core + # google-cloud-core + # google-cloud-storage + # runai-model-streamer-gcs +google-cloud-core==2.4.3 + # via google-cloud-storage +google-cloud-storage==3.4.0 + # via runai-model-streamer-gcs +google-crc32c==1.7.1 + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.7.2 + # via google-cloud-storage +googleapis-common-protos==1.70.0 + # via + # google-api-core + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +gpt-oss==0.0.8 + # via -r requirements/test/cuda.in +graphql-core==3.2.6 + # via hypothesis-graphql +greenlet==3.2.3 + # via sqlalchemy +grpcio==1.78.0 + # via + # -r requirements/test/cuda.in + # grpcio-reflection + # opentelemetry-exporter-otlp-proto-grpc + # ray +grpcio-reflection==1.78.0 + # via -r requirements/test/cuda.in +h11==0.14.0 + # via + # httpcore + # uvicorn +h2==4.3.0 + # via httpx +harfile==0.5.0 + # via schemathesis +hf-xet==1.5.1 + # via huggingface-hub +hiredis==3.0.0 + # via tensorizer +hpack==4.1.0 + # via h2 +html2text==2025.4.15 + # via gpt-oss +httpcore==1.0.6 + # via httpx +httptools==0.8.0 + # via uvicorn +httpx==0.27.2 + # via + # -r requirements/test/cuda.in + # anthropic + # fastapi + # fastapi-cloud-cli + # huggingface-hub + # mcp + # model-hosting-container-standards + # openai + # perceptron + # schemathesis +httpx-sse==0.4.3 + # via mcp +huggingface-hub==1.22.0 + # via + # accelerate + # datasets + # evaluate + # open-clip-torch + # peft + # segmentation-models-pytorch + # sentence-transformers + # timm + # tokenizers + # transformers + # vocos +humanize==4.11.0 + # via runai-model-streamer +hyperframe==6.1.0 + # via h2 +hypothesis==6.131.0 + # via + # hypothesis-graphql + # hypothesis-jsonschema + # schemathesis +hypothesis-graphql==0.13.0 + # via schemathesis +hypothesis-jsonschema==0.23.1 + # via schemathesis +idna==3.10 + # via + # anyio + # email-validator + # httpx + # requests + # yarl +ijson==3.5.0 + # via -r requirements/test/../common.txt +imagehash==4.3.2 + # via -r requirements/test/cuda.in +imageio==2.37.0 + # via scikit-image +importlib-metadata==8.7.0 + # via opentelemetry-api +inflect==5.6.2 + # via datamodel-code-generator +iniconfig==2.0.0 + # via pytest +instanttensor==0.1.5 + # via -r requirements/test/cuda.in +interegular==0.3.3 + # via lm-format-enforcer +isodate==0.7.2 + # via azure-storage-blob +isort==5.13.2 + # via datamodel-code-generator +jinja2==3.1.6 + # via + # datamodel-code-generator + # fastapi + # genai-perf + # lm-eval + # torch +jiter==0.15.0 + # via + # anthropic + # openai +jiwer==3.0.5 + # via -r requirements/test/cuda.in +jmespath==1.0.1 + # via + # boto3 + # botocore + # model-hosting-container-standards +joblib==1.4.2 + # via + # librosa + # nltk + # scikit-learn +jsonschema==4.23.0 + # via + # -r requirements/test/../common.txt + # hypothesis-jsonschema + # mcp + # mistral-common + # ray +jsonschema-rs==0.46.5 + # via schemathesis +jsonschema-specifications==2024.10.1 + # via jsonschema +junit-xml==1.9 + # via schemathesis +kaldi-native-fbank==1.22.3 + # via -r requirements/test/cuda.in +kaleido==0.2.1 + # via genai-perf +kiwisolver==1.4.7 + # via matplotlib +lark==1.2.2 + # via -r requirements/test/../common.txt +lazy-loader==0.4 + # via + # librosa + # scikit-image +libnacl==2.1.0 + # via tensorizer +librosa==0.10.2.post1 + # via -r requirements/test/cuda.in +llguidance==1.7.6 + # via -r requirements/test/../common.txt +llvmlite==0.47.0 + # via numba +lm-eval==0.4.12 + # via -r requirements/test/cuda.in +lm-format-enforcer==0.11.3 + # via -r requirements/test/../common.txt +loguru==0.7.3 + # via compressed-tensors +lxml==5.3.0 + # via + # blobfile + # gpt-oss + # sacrebleu +mako==1.3.10 + # via alembic +markdown-it-py==3.0.0 + # via rich +markupsafe==3.0.1 + # via + # jinja2 + # mako + # werkzeug +matplotlib==3.9.2 + # via -r requirements/test/cuda.in +mbstrdecoder==1.1.3 + # via + # dataproperty + # pytablewriter + # typepy +mcp==1.28.1 + # via -r requirements/test/../common.txt +mdurl==0.1.2 + # via markdown-it-py +mistral-common==1.11.5 + # via + # -r requirements/test/../common.txt + # -r requirements/test/cuda.in +model-hosting-container-standards==0.1.16 + # via -r requirements/test/../common.txt +more-itertools==10.5.0 + # via lm-eval +mpmath==1.3.0 + # via sympy +msal==1.34.0 + # via + # azure-identity + # msal-extensions +msal-extensions==1.3.1 + # via azure-identity +msgpack==1.1.0 + # via + # librosa + # ray +msgspec==0.21.1 + # via -r requirements/test/../common.txt +mteb==2.8.3 + # via -r requirements/test/cuda.in +multidict==6.1.0 + # via + # aiohttp + # yarl +multiprocess==0.70.16 + # via + # datasets + # evaluate +mypy-extensions==1.0.0 + # via black +networkx==3.2.1 + # via + # scikit-image + # torch +ninja==1.13.0 + # via -r requirements/test/../common.txt +nltk==3.9.1 + # via rouge-score +num2words==0.5.14 + # via -r requirements/test/cuda.in +numba==0.65.0 + # via + # -r requirements/test/cuda.in + # librosa +numpy==2.2.6 + # via + # -r requirements/test/../common.txt + # accelerate + # albumentations + # bitsandbytes + # bm25s + # contourpy + # cupy-cuda12x + # datasets + # decord + # einx + # encodec + # evaluate + # fastparquet + # genai-perf + # imagehash + # imageio + # librosa + # lm-eval + # matplotlib + # mistral-common + # mteb + # numba + # opencv-python-headless + # optuna + # pandas + # patsy + # peft + # perceptron + # pywavelets + # rouge-score + # runai-model-streamer + # sacrebleu + # scikit-image + # scikit-learn + # scipy + # segmentation-models-pytorch + # soxr + # statsmodels + # tensorizer + # tifffile + # torchvision + # transformers + # tritonclient + # vocos + # xgrammar +open-clip-torch==2.32.0 + # via -r requirements/test/cuda.in +openai==2.44.0 + # via -r requirements/test/../common.txt +openai-harmony==0.0.4 + # via + # -r requirements/test/../common.txt + # gpt-oss +opencensus==0.11.4 + # via ray +opencensus-context==0.1.3 + # via opencensus +opencv-python-headless==4.13.0.90 + # via + # -r requirements/test/../common.txt + # albumentations + # mistral-common +openpyxl==3.1.5 + # via -r requirements/test/cuda.in +opentelemetry-api==1.35.0 + # via + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-exporter-prometheus + # opentelemetry-sdk + # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp==1.35.0 + # via -r requirements/test/../common.txt +opentelemetry-exporter-otlp-proto-common==1.35.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-grpc==1.35.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-otlp-proto-http==1.35.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-prometheus==0.56b0 + # via ray +opentelemetry-proto==1.35.0 + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # ray +opentelemetry-sdk==1.35.0 + # via + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-exporter-prometheus + # ray +opentelemetry-semantic-conventions==0.56b0 + # via opentelemetry-sdk +opentelemetry-semantic-conventions-ai==0.4.13 + # via -r requirements/test/../common.txt +optuna==3.6.1 + # via genai-perf +orjson==3.11.5 + # via genai-perf +outlines-core==0.2.14 + # via -r requirements/test/../common.txt +packaging==24.2 + # via + # accelerate + # bitsandbytes + # black + # datamodel-code-generator + # datasets + # evaluate + # fastparquet + # huggingface-hub + # lazy-loader + # lm-format-enforcer + # matplotlib + # optuna + # peft + # plotly + # pooch + # pytest + # pytest-rerunfailures + # ray + # scikit-image + # statsmodels + # transformers + # typepy +pandas==2.2.3 + # via + # datasets + # evaluate + # fastparquet + # genai-perf + # statsmodels +partial-json-parser==0.2.1.1.post7 + # via -r requirements/test/../common.txt +pathspec==0.12.1 + # via black +pathvalidate==3.2.1 + # via pytablewriter +patsy==1.0.1 + # via statsmodels +peft==0.19.1 + # via -r requirements/test/cuda.in +perceptron==0.1.4 + # via -r requirements/test/cuda.in +perf-analyzer==0.1.0 + # via genai-perf +pillow==10.4.0 + # via + # -r requirements/test/../common.txt + # genai-perf + # imagehash + # imageio + # matplotlib + # mistral-common + # perceptron + # scikit-image + # segmentation-models-pytorch + # torchvision +platformdirs==4.3.6 + # via + # black + # pooch + # virtualenv +plotly==5.24.1 + # via + # -r requirements/test/cuda.in + # genai-perf +pluggy==1.5.0 + # via + # pytest + # pytest-cov +polars==1.29.0 + # via mteb +pooch==1.8.2 + # via librosa +portalocker==2.10.1 + # via sacrebleu +pqdm==0.2.0 + # via -r requirements/test/cuda.in +prometheus-client==0.22.0 + # via + # -r requirements/test/../common.txt + # opentelemetry-exporter-prometheus + # prometheus-fastapi-instrumentator + # ray +prometheus-fastapi-instrumentator==8.0.2 + # via -r requirements/test/../common.txt +propcache==0.2.0 + # via + # aiohttp + # yarl +proto-plus==1.26.1 + # via google-api-core +protobuf==6.33.6 + # via + # -r requirements/test/../common.txt + # google-api-core + # googleapis-common-protos + # grpcio-reflection + # opentelemetry-proto + # proto-plus + # ray + # tensorizer +psutil==6.1.0 + # via + # -r requirements/test/../common.txt + # accelerate + # peft + # tensorizer +py==1.11.0 + # via pytest-forked +py-cpuinfo==9.0.0 + # via -r requirements/test/../common.txt +py-spy==0.4.0 + # via ray +pyarrow==23.0.0 + # via + # datasets + # genai-perf +pyasn1==0.6.1 + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.4.2 + # via google-auth +pybase64==1.4.3 + # via -r requirements/test/../common.txt +pycountry==24.6.1 + # via pydantic-extra-types +pycparser==2.22 + # via cffi +pycryptodomex==3.22.0 + # via blobfile +pydantic==2.12.0 + # via + # -r requirements/test/../common.txt + # albumentations + # anthropic + # compressed-tensors + # datamodel-code-generator + # fastapi + # fastapi-cloud-cli + # gpt-oss + # lm-format-enforcer + # mcp + # mistral-common + # model-hosting-container-standards + # mteb + # openai + # openai-harmony + # pydantic-extra-types + # pydantic-settings + # ray + # xgrammar +pydantic-core==2.41.1 + # via pydantic +pydantic-extra-types==2.10.5 + # via + # fastapi + # mistral-common +pydantic-settings==2.14.2 + # via + # fastapi + # mcp +pygments==2.18.0 + # via + # pytest + # rich +pyjwt==2.11.0 + # via + # mcp + # msal +pyparsing==3.2.0 + # via matplotlib +pyrate-limiter==4.4.0 + # via schemathesis +pystemmer==3.0.0 + # via mteb +pytablewriter==1.2.0 + # via lm-eval +pytest==9.1.0 + # via + # -r requirements/test/cuda.in + # buildkite-test-collector + # genai-perf + # pytest-asyncio + # pytest-cov + # pytest-forked + # pytest-mock + # pytest-rerunfailures + # pytest-shard + # pytest-timeout + # schemathesis +pytest-asyncio==1.4.0 + # via -r requirements/test/cuda.in +pytest-cov==6.3.0 + # via -r requirements/test/cuda.in +pytest-forked==1.6.0 + # via -r requirements/test/cuda.in +pytest-mock==3.14.0 + # via genai-perf +pytest-rerunfailures==14.0 + # via -r requirements/test/cuda.in +pytest-shard==0.1.2 + # via -r requirements/test/cuda.in +pytest-timeout==2.3.1 + # via -r requirements/test/cuda.in +python-dateutil==2.9.0.post0 + # via + # botocore + # matplotlib + # pandas + # typepy +python-dotenv==1.2.2 + # via + # pydantic-settings + # uvicorn +python-json-logger==4.1.0 + # via -r requirements/test/../common.txt +python-multipart==0.0.32 + # via + # fastapi + # mcp +python-rapidjson==1.20 + # via tritonclient +pytrec-eval-terrier==0.5.7 + # via mteb +pytz==2024.2 + # via + # pandas + # typepy +pywavelets==1.9.0 + # via imagehash +pyyaml==6.0.2 + # via + # -r requirements/test/../common.txt + # accelerate + # albumentations + # datamodel-code-generator + # datasets + # genai-perf + # huggingface-hub + # lm-format-enforcer + # optuna + # peft + # ray + # responses + # schemathesis + # timm + # transformers + # uvicorn + # vocos +pyzmq==27.1.0 + # via -r requirements/test/../common.txt +rapidfuzz==3.12.1 + # via jiwer +ray==2.48.0 + # via -r requirements/test/cuda.in +redis==5.2.0 + # via tensorizer +referencing==0.35.1 + # via + # jsonschema + # jsonschema-specifications +regex==2026.2.28 + # via + # -r requirements/test/../common.txt + # nltk + # open-clip-torch + # sacrebleu + # tiktoken + # transformers +requests==2.32.3 + # via + # -r requirements/test/../common.txt + # azure-core + # buildkite-test-collector + # datasets + # docker + # evaluate + # google-api-core + # google-cloud-storage + # gpt-oss + # lm-eval + # mistral-common + # msal + # mteb + # opentelemetry-exporter-otlp-proto-http + # pooch + # ray + # responses + # schemathesis + # starlette-testclient + # tiktoken +responses==0.25.3 + # via genai-perf +rich==13.9.4 + # via + # genai-perf + # mteb + # perceptron + # rich-toolkit + # schemathesis + # typer +rich-toolkit==0.20.1 + # via + # fastapi-cli + # fastapi-cloud-cli +rignore==0.7.6 + # via fastapi-cloud-cli +rouge-score==0.1.2 + # via lm-eval +rpds-py==0.20.1 + # via + # jsonschema + # referencing +rsa==4.9.1 + # via google-auth +runai-model-streamer==0.15.7 + # via -r requirements/test/cuda.in +runai-model-streamer-azure==0.15.7 + # via runai-model-streamer +runai-model-streamer-gcs==0.15.7 + # via runai-model-streamer +runai-model-streamer-s3==0.15.7 + # via runai-model-streamer +s3transfer==0.10.3 + # via boto3 +sacrebleu==2.4.3 + # via lm-eval +safetensors==0.8.0 + # via + # -r requirements/test/../common.txt + # accelerate + # open-clip-torch + # peft + # segmentation-models-pytorch + # timm + # transformers +schemathesis==4.21.6 + # via -r requirements/test/cuda.in +scikit-image==0.25.2 + # via albumentations +scikit-learn==1.5.2 + # via + # albumentations + # librosa + # lm-eval + # mteb + # sentence-transformers +scipy==1.13.1 + # via + # albumentations + # bm25s + # imagehash + # librosa + # mteb + # scikit-image + # scikit-learn + # sentence-transformers + # statsmodels + # vocos +segmentation-models-pytorch==0.5.0 + # via -r requirements/test/cuda.in +sentence-transformers==5.2.0 + # via + # -r requirements/test/cuda.in + # mteb +sentencepiece==0.2.1 + # via -r requirements/test/../common.txt +sentry-sdk==2.63.0 + # via fastapi-cloud-cli +setproctitle==1.3.7 + # via -r requirements/test/../common.txt +setuptools==77.0.3 + # via + # -r requirements/test/../common.txt + # model-hosting-container-standards + # pytablewriter + # torch +shellingham==1.5.4 + # via + # perceptron + # typer +six==1.16.0 + # via + # -r requirements/test/../common.txt + # junit-xml + # opencensus + # python-dateutil + # rouge-score +smart-open==7.1.0 + # via ray +sniffio==1.3.1 + # via + # anthropic + # httpx + # openai +sortedcontainers==2.4.0 + # via hypothesis +soundfile==0.12.1 + # via + # -r requirements/test/cuda.in + # genai-perf + # librosa + # mistral-common +soxr==0.5.0.post1 + # via + # librosa + # mistral-common +sqlalchemy==2.0.41 + # via + # alembic + # optuna +sqlitedict==2.1.0 + # via lm-eval +sse-starlette==3.4.5 + # via mcp +starlette==1.3.1 + # via + # -r requirements/test/../common.txt + # fastapi + # mcp + # model-hosting-container-standards + # prometheus-fastapi-instrumentator + # sse-starlette + # starlette-testclient +starlette-testclient==0.4.1 + # via schemathesis +statsmodels==0.14.4 + # via genai-perf +structlog==25.4.0 + # via gpt-oss +supervisor==4.3.0 + # via model-hosting-container-standards +sympy==1.13.3 + # via + # einx + # torch +tabledata==1.3.3 + # via pytablewriter +tabulate==0.9.0 + # via sacrebleu +tblib==3.1.0 + # via -r requirements/test/cuda.in +tcolorpy==0.1.6 + # via pytablewriter +tenacity==9.1.2 + # via + # gpt-oss + # lm-eval + # plotly + # schemathesis +tensorizer==2.10.1 + # via -r requirements/test/cuda.in +termcolor==3.1.0 + # via gpt-oss +threadpoolctl==3.5.0 + # via scikit-learn +tifffile==2025.3.30 + # via scikit-image +tiktoken==0.12.0 + # via + # -r requirements/test/../common.txt + # gpt-oss + # lm-eval + # mistral-common +timm==1.0.17 + # via + # -r requirements/test/cuda.in + # open-clip-torch + # segmentation-models-pytorch +tokenizers==0.22.2 + # via + # -r requirements/test/../common.txt + # -r requirements/test/cuda.in + # transformers +torch==2.11.0+cpu + # via + # -r requirements/test/cuda.in + # accelerate + # bitsandbytes + # compressed-tensors + # encodec + # instanttensor + # mteb + # open-clip-torch + # peft + # runai-model-streamer + # segmentation-models-pytorch + # sentence-transformers + # tensorizer + # timm + # torchvision + # vector-quantize-pytorch + # vocos + # xgrammar +torchaudio==2.11.0+cpu + # via + # -r requirements/test/cuda.in + # encodec + # vocos +torchcodec==0.14.0+cpu + # via -r requirements/test/cuda.in +torchvision==0.26.0+cpu + # via + # -r requirements/test/cuda.in + # open-clip-torch + # segmentation-models-pytorch + # timm +tqdm==4.67.3 + # via + # -r requirements/test/../common.txt + # datasets + # evaluate + # huggingface-hub + # lm-eval + # mteb + # nltk + # open-clip-torch + # openai + # optuna + # peft + # pqdm + # segmentation-models-pytorch + # sentence-transformers + # transformers +transformers==5.13.1 + # via + # -r requirements/test/../common.txt + # -r requirements/test/cuda.in + # compressed-tensors + # genai-perf + # peft + # sentence-transformers + # transformers-stream-generator + # xgrammar +transformers-stream-generator==0.0.5 + # via -r requirements/test/cuda.in +triton==3.6.0 + # via xgrammar +tritonclient==2.64.0 + # via -r requirements/test/cuda.in +typepy==1.3.2 + # via + # dataproperty + # pytablewriter + # tabledata +typer==0.26.8 + # via + # fastapi-cli + # fastapi-cloud-cli + # fastsafetensors + # perceptron + # transformers +typing-extensions==4.15.0 + # via + # -r requirements/test/../common.txt + # aiosignal + # albumentations + # alembic + # anthropic + # anyio + # apache-tvm-ffi + # azure-core + # azure-identity + # azure-storage-blob + # chz + # fastapi + # grpcio + # huggingface-hub + # librosa + # lm-eval + # mcp + # mistral-common + # mteb + # openai + # opentelemetry-api + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions + # pqdm + # pydantic + # pydantic-core + # pydantic-extra-types + # pytest-asyncio + # rich-toolkit + # schemathesis + # sentence-transformers + # sqlalchemy + # starlette + # torch + # typing-inspection + # xgrammar +typing-inspection==0.4.2 + # via + # fastapi + # mcp + # pydantic + # pydantic-settings +tzdata==2024.2 + # via pandas +urllib3==2.2.3 + # via + # blobfile + # botocore + # docker + # requests + # responses + # sentry-sdk + # tritonclient +uvicorn==0.35.0 + # via + # fastapi + # fastapi-cli + # fastapi-cloud-cli + # gpt-oss + # mcp +uvloop==0.22.1 + # via uvicorn +vector-quantize-pytorch==1.21.2 + # via -r requirements/test/cuda.in +virtualenv==20.31.2 + # via ray +vocos==0.1.0 + # via -r requirements/test/cuda.in +watchfiles==1.2.0 + # via + # -r requirements/test/../common.txt + # uvicorn +wcwidth==0.2.13 + # via ftfy +websockets==16.0 + # via uvicorn +werkzeug==3.1.3 + # via schemathesis +word2number==1.1 + # via lm-eval +wrapt==1.17.2 + # via smart-open +xgrammar==0.2.3 + # via -r requirements/test/../common.txt +xxhash==3.5.0 + # via + # datasets + # evaluate +yarl==1.17.1 + # via aiohttp +zipp==3.23.0 + # via importlib-metadata diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 8d7ad7d0aa2..bd6f179c105 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -1,3 +1,5 @@ +-r ../common.txt + # testing pytest tensorizer==2.10.1 @@ -11,14 +13,14 @@ pytest-cov # testing utils albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests +torchcodec >= 0.14 # required for torchcodec video backend tests backoff # required for phi4mm test blobfile # required for kimi-vl test -einops # required for MPT, qwen-vl httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.18.1 # required for phi-4-mm test +peft>=0.19.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -31,16 +33,15 @@ torchaudio==2.11.0 torchvision==0.26.0 transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.3 # required for voxtral test +mistral_common[image,audio] >= 1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py -opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.13.1 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test. +schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 @@ -55,11 +56,9 @@ grpcio-reflection==1.78.0 arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test numba == 0.65.0 # Required for N-gram speculative decoding -numpy runai-model-streamer[s3,gcs,azure]==0.15.7 fastsafetensors>=0.3.2 instanttensor>=0.1.5; platform_machine == "x86_64" -pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0; platform_machine == "x86_64" # terratorch is temporarily disabled while PyPI has the `lightning` package # in `quarantined` status (every published terratorch version transitively @@ -73,6 +72,7 @@ gpt-oss >= 0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test +cohere_melody>=0.9.0 # required for cohere command reasoning parser test # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with teerratorch requirements. diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index a3e1466c763..94986ec2479 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -9,6 +9,7 @@ aiohappyeyeballs==2.6.1 aiohttp==3.13.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiohttp-cors # datasets # fsspec @@ -24,25 +25,39 @@ albumentations==1.4.6 alembic==1.16.4 # via optuna annotated-doc==0.0.4 - # via fastapi + # via + # fastapi + # typer annotated-types==0.7.0 # via pydantic -anyio==4.6.2.post1 +anthropic==0.112.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +anyio==4.14.1 + # via + # anthropic # httpx + # mcp + # openai + # sse-starlette # starlette + # watchfiles +apache-tvm-ffi==0.1.9 + # via + # -c requirements/cuda.txt + # xgrammar arctic-inference==0.1.1 # via -r requirements/test/cuda.in argcomplete==3.5.1 # via datamodel-code-generator -arrow==1.3.0 - # via isoduration +astor==0.8.1 + # via depyf attrs==24.2.0 # via # aiohttp # hypothesis # jsonschema - # pytest-subtests # referencing audioread==3.0.1 # via librosa @@ -57,13 +72,13 @@ azure-identity==1.25.2 azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 - # via - # -r requirements/test/cuda.in - # schemathesis + # via -r requirements/test/cuda.in bitsandbytes==0.49.2 # via -r requirements/test/cuda.in black==24.10.0 # via datamodel-code-generator +blake3==1.0.9 + # via -r requirements/test/../common.txt blobfile==3.0.0 # via -r requirements/test/cuda.in bm25s==0.2.13 @@ -81,12 +96,17 @@ bounded-pool-executor==0.0.3 buildkite-test-collector==0.1.9 # via -r requirements/test/cuda.in cachetools==5.5.2 - # via google-auth + # via + # -r requirements/test/../common.txt + # google-auth +cbor2==6.1.2 + # via -r requirements/test/../common.txt certifi==2024.8.30 # via # httpcore # httpx # requests + # sentry-sdk cffi==2.0.0 # via # cryptography @@ -97,24 +117,32 @@ charset-normalizer==3.4.0 # via requests chz==0.3.0 # via gpt-oss -click==8.1.7 +click==8.4.2 # via # black + # huggingface-hub # jiwer # nltk # ray + # rich-toolkit # schemathesis - # typer # uvicorn +cloudpickle==3.1.2 + # via -r requirements/test/../common.txt +cohere-melody==0.9.0 + # via -r requirements/test/cuda.in colorama==0.4.6 # via # perceptron # sacrebleu - # schemathesis colorful==0.5.6 # via ray colorlog==6.10.1 # via optuna +compressed-tensors==0.17.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt contourpy==1.3.0 # via matplotlib coverage==7.10.6 @@ -153,40 +181,69 @@ decorator==5.1.1 # via librosa decord==0.6.0 # via -r requirements/test/cuda.in +depyf==0.20.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +detect-installer==0.1.0 + # via fastapi-cloud-cli dill==0.3.8 # via # datasets + # depyf # evaluate # lm-eval # multiprocess +diskcache==5.6.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt distlib==0.3.9 # via virtualenv +distro==1.9.0 + # via + # anthropic + # openai dnspython==2.7.0 # via email-validator docker==7.1.0 # via gpt-oss docopt==0.6.2 # via num2words +docstring-parser==0.18.0 + # via anthropic einops==0.8.1 # via - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # encodec # vector-quantize-pytorch # vocos einx==0.3.0 # via vector-quantize-pytorch email-validator==2.2.0 - # via pydantic + # via + # fastapi + # pydantic encodec==0.1.1 # via vocos et-xmlfile==2.0.0 # via openpyxl evaluate==0.4.3 # via lm-eval -fastapi==0.128.0 +fastapi==0.136.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss + # model-hosting-container-standards +fastapi-cli==0.0.27 + # via fastapi +fastapi-cloud-cli==0.21.0 + # via fastapi-cli +fastar==0.11.0 + # via + # fastapi + # fastapi-cloud-cli fastparquet==2024.11.0 # via genai-perf fastrlock==0.8.2 @@ -198,6 +255,7 @@ fastsafetensors==0.3.2 filelock==3.16.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # blobfile # datasets # huggingface-hub @@ -206,8 +264,6 @@ filelock==3.16.1 # virtualenv fonttools==4.55.0 # via matplotlib -fqdn==1.5.1 - # via jsonschema frozendict==2.4.6 # via einx frozenlist==1.5.0 @@ -249,7 +305,10 @@ google-crc32c==1.7.1 google-resumable-media==2.7.2 # via google-cloud-storage googleapis-common-protos==1.70.0 - # via google-api-core + # via + # google-api-core + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http gpt-oss==0.0.8 # via -r requirements/test/cuda.in graphql-core==3.2.6 @@ -260,6 +319,7 @@ grpcio==1.78.0 # via # -r requirements/test/cuda.in # grpcio-reflection + # opentelemetry-exporter-otlp-proto-grpc # ray grpcio-reflection==1.78.0 # via -r requirements/test/cuda.in @@ -269,9 +329,9 @@ h11==0.14.0 # uvicorn h2==4.3.0 # via httpx -harfile==0.3.0 +harfile==0.5.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub hiredis==3.0.0 # via tensorizer @@ -281,13 +341,23 @@ html2text==2025.4.15 # via gpt-oss httpcore==1.0.6 # via httpx +httptools==0.8.0 + # via uvicorn httpx==0.27.2 # via # -r requirements/test/cuda.in + # anthropic + # fastapi + # fastapi-cloud-cli # huggingface-hub + # mcp + # model-hosting-container-standards + # openai # perceptron # schemathesis -huggingface-hub==1.10.2 +httpx-sse==0.4.3 + # via mcp +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -309,7 +379,7 @@ hypothesis==6.131.0 # hypothesis-graphql # hypothesis-jsonschema # schemathesis -hypothesis-graphql==0.11.1 +hypothesis-graphql==0.13.0 # via schemathesis hypothesis-jsonschema==0.23.1 # via schemathesis @@ -318,9 +388,10 @@ idna==3.10 # anyio # email-validator # httpx - # jsonschema # requests # yarl +ijson==3.5.0 + # via -r requirements/test/../common.txt imagehash==4.3.2 # via -r requirements/test/cuda.in imageio==2.37.0 @@ -333,37 +404,45 @@ iniconfig==2.0.0 # via pytest instanttensor==0.1.5 # via -r requirements/test/cuda.in +interegular==0.3.3 + # via lm-format-enforcer isodate==0.7.2 # via azure-storage-blob -isoduration==20.11.0 - # via jsonschema isort==5.13.2 # via datamodel-code-generator jinja2==3.1.6 # via # datamodel-code-generator + # fastapi # genai-perf # lm-eval # torch +jiter==0.15.0 + # via + # anthropic + # openai jiwer==3.0.5 # via -r requirements/test/cuda.in jmespath==1.0.1 # via # boto3 # botocore + # model-hosting-container-standards joblib==1.4.2 # via # librosa # nltk # scikit-learn -jsonpointer==3.0.0 - # via jsonschema jsonschema==4.23.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # hypothesis-jsonschema + # mcp # mistral-common # ray - # schemathesis +jsonschema-rs==0.46.5 + # via schemathesis jsonschema-specifications==2024.10.1 # via jsonschema junit-xml==1.9 @@ -374,6 +453,10 @@ kaleido==0.2.1 # via genai-perf kiwisolver==1.4.7 # via matplotlib +lark==1.2.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt lazy-loader==0.4 # via # librosa @@ -382,10 +465,20 @@ libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 # via -r requirements/test/cuda.in +llguidance==1.7.6 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt llvmlite==0.47.0 # via numba lm-eval==0.4.12 # via -r requirements/test/cuda.in +lm-format-enforcer==0.11.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +loguru==0.7.3 + # via compressed-tensors lxml==5.3.0 # via # blobfile @@ -407,12 +500,19 @@ mbstrdecoder==1.1.3 # dataproperty # pytablewriter # typepy +mcp==1.28.1 + # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in +model-hosting-container-standards==0.1.16 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt more-itertools==10.5.0 # via lm-eval mpmath==1.3.0 @@ -427,6 +527,8 @@ msgpack==1.1.0 # via # librosa # ray +msgspec==0.21.1 + # via -r requirements/test/../common.txt mteb==2.8.3 # via -r requirements/test/cuda.in multidict==6.1.0 @@ -443,6 +545,8 @@ networkx==3.2.1 # via # scikit-image # torch +ninja==1.13.0 + # via -r requirements/test/../common.txt nltk==3.9.1 # via rouge-score num2words==0.5.14 @@ -454,7 +558,7 @@ numba==0.65.0 # librosa numpy==2.2.6 # via - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # accelerate # albumentations # bitsandbytes @@ -498,6 +602,7 @@ numpy==2.2.6 # transformers # tritonclient # vocos + # xgrammar nvidia-cublas==13.1.0.3 # via # cuda-toolkit @@ -539,9 +644,14 @@ nvidia-nvtx==13.0.85 # via cuda-toolkit open-clip-torch==2.32.0 # via -r requirements/test/cuda.in +openai==2.44.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt openai-harmony==0.0.4 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss opencensus==0.11.4 # via ray @@ -550,7 +660,7 @@ opencensus-context==0.1.3 opencv-python-headless==4.13.0.90 # via # -c requirements/common.txt - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # albumentations # mistral-common openpyxl==3.1.5 @@ -558,24 +668,54 @@ openpyxl==3.1.5 opentelemetry-api==1.35.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp==1.35.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +opentelemetry-exporter-otlp-proto-common==1.35.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-grpc==1.35.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-otlp-proto-http==1.35.0 + # via opentelemetry-exporter-otlp opentelemetry-exporter-prometheus==0.56b0 # via ray opentelemetry-proto==1.35.0 - # via ray + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # ray opentelemetry-sdk==1.35.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus # ray opentelemetry-semantic-conventions==0.56b0 # via opentelemetry-sdk +opentelemetry-semantic-conventions-ai==0.4.13 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt optuna==3.6.1 # via genai-perf orjson==3.11.5 # via genai-perf +outlines-core==0.2.14 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt packaging==24.2 # via # accelerate @@ -587,6 +727,7 @@ packaging==24.2 # fastparquet # huggingface-hub # lazy-loader + # lm-format-enforcer # matplotlib # optuna # peft @@ -606,13 +747,15 @@ pandas==2.2.3 # fastparquet # genai-perf # statsmodels +partial-json-parser==0.2.1.1.post7 + # via -r requirements/test/../common.txt pathspec==0.12.1 # via black pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.18.1 +peft==0.19.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in @@ -620,6 +763,7 @@ perf-analyzer==0.1.0 # via genai-perf pillow==10.4.0 # via + # -r requirements/test/../common.txt # genai-perf # imagehash # imageio @@ -653,8 +797,14 @@ pqdm==0.2.0 prometheus-client==0.22.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # opentelemetry-exporter-prometheus + # prometheus-fastapi-instrumentator # ray +prometheus-fastapi-instrumentator==8.0.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt propcache==0.2.0 # via # aiohttp @@ -664,6 +814,7 @@ proto-plus==1.26.1 protobuf==6.33.6 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # google-api-core # googleapis-common-protos # grpcio-reflection @@ -673,11 +824,14 @@ protobuf==6.33.6 # tensorizer psutil==6.1.0 # via + # -r requirements/test/../common.txt # accelerate # peft # tensorizer py==1.11.0 # via pytest-forked +py-cpuinfo==9.0.0 + # via -r requirements/test/../common.txt py-spy==0.4.0 # via ray pyarrow==23.0.0 @@ -690,6 +844,8 @@ pyasn1==0.6.1 # rsa pyasn1-modules==0.4.2 # via google-auth +pybase64==1.4.3 + # via -r requirements/test/../common.txt pycountry==24.6.1 # via pydantic-extra-types pycparser==2.22 @@ -699,33 +855,52 @@ pycryptodomex==3.22.0 pydantic==2.12.0 # via # -c requirements/common.txt - # -r requirements/test/cuda.in + # -r requirements/test/../common.txt # albumentations + # anthropic + # compressed-tensors # datamodel-code-generator # fastapi + # fastapi-cloud-cli # gpt-oss + # lm-format-enforcer + # mcp # mistral-common + # model-hosting-container-standards # mteb + # openai # openai-harmony # pydantic-extra-types + # pydantic-settings # ray + # xgrammar pydantic-core==2.41.1 # via pydantic pydantic-extra-types==2.10.5 - # via mistral-common + # via + # fastapi + # mistral-common +pydantic-settings==2.14.2 + # via + # fastapi + # mcp pygments==2.18.0 - # via rich + # via + # pytest + # rich pyjwt==2.11.0 - # via msal + # via + # mcp + # msal pyparsing==3.2.0 # via matplotlib -pyrate-limiter==3.7.0 +pyrate-limiter==4.4.0 # via schemathesis pystemmer==3.0.0 # via mteb pytablewriter==1.2.0 # via lm-eval -pytest==8.3.5 +pytest==9.1.0 # via # -r requirements/test/cuda.in # buildkite-test-collector @@ -736,10 +911,9 @@ pytest==8.3.5 # pytest-mock # pytest-rerunfailures # pytest-shard - # pytest-subtests # pytest-timeout # schemathesis -pytest-asyncio==0.24.0 +pytest-asyncio==1.4.0 # via -r requirements/test/cuda.in pytest-cov==6.3.0 # via -r requirements/test/cuda.in @@ -751,17 +925,24 @@ pytest-rerunfailures==14.0 # via -r requirements/test/cuda.in pytest-shard==0.1.2 # via -r requirements/test/cuda.in -pytest-subtests==0.14.1 - # via schemathesis pytest-timeout==2.3.1 # via -r requirements/test/cuda.in python-dateutil==2.9.0.post0 # via - # arrow # botocore # matplotlib # pandas # typepy +python-dotenv==1.2.2 + # via + # pydantic-settings + # uvicorn +python-json-logger==4.1.0 + # via -r requirements/test/../common.txt +python-multipart==0.0.32 + # via + # fastapi + # mcp python-rapidjson==1.20 # via tritonclient pytrec-eval-terrier==0.5.7 @@ -774,12 +955,14 @@ pywavelets==1.9.0 # via imagehash pyyaml==6.0.2 # via + # -r requirements/test/../common.txt # accelerate # albumentations # datamodel-code-generator # datasets # genai-perf # huggingface-hub + # lm-format-enforcer # optuna # peft # ray @@ -787,7 +970,12 @@ pyyaml==6.0.2 # schemathesis # timm # transformers + # uvicorn # vocos +pyzmq==27.1.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt rapidfuzz==3.12.1 # via jiwer ray==2.48.0 @@ -800,6 +988,7 @@ referencing==0.35.1 # jsonschema-specifications regex==2026.2.28 # via + # -r requirements/test/../common.txt # nltk # open-clip-torch # sacrebleu @@ -808,6 +997,7 @@ regex==2026.2.28 requests==2.32.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # azure-core # buildkite-test-collector # datasets @@ -820,6 +1010,7 @@ requests==2.32.3 # mistral-common # msal # mteb + # opentelemetry-exporter-otlp-proto-http # pooch # ray # responses @@ -828,16 +1019,20 @@ requests==2.32.3 # tiktoken responses==0.25.3 # via genai-perf -rfc3339-validator==0.1.4 - # via jsonschema -rfc3987==1.3.8 - # via jsonschema rich==13.9.4 # via # genai-perf # mteb # perceptron + # rich-toolkit + # schemathesis # typer +rich-toolkit==0.20.1 + # via + # fastapi-cli + # fastapi-cloud-cli +rignore==0.7.6 + # via fastapi-cloud-cli rouge-score==0.1.2 # via lm-eval rpds-py==0.20.1 @@ -858,16 +1053,17 @@ s3transfer==0.10.3 # via boto3 sacrebleu==2.4.3 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # accelerate # open-clip-torch # peft # segmentation-models-pytorch # timm # transformers -schemathesis==3.39.15 +schemathesis==4.21.6 # via -r requirements/test/cuda.in scikit-image==0.25.2 # via albumentations @@ -896,9 +1092,17 @@ sentence-transformers==5.2.0 # via # -r requirements/test/cuda.in # mteb +sentencepiece==0.2.1 + # via -r requirements/test/../common.txt +sentry-sdk==2.63.0 + # via fastapi-cloud-cli +setproctitle==1.3.7 + # via -r requirements/test/../common.txt setuptools==77.0.3 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # model-hosting-container-standards # pytablewriter # torch shellingham==1.5.4 @@ -908,17 +1112,18 @@ shellingham==1.5.4 six==1.16.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # junit-xml # opencensus # python-dateutil - # rfc3339-validator # rouge-score smart-open==7.1.0 # via ray sniffio==1.3.1 # via - # anyio + # anthropic # httpx + # openai sortedcontainers==2.4.0 # via hypothesis soundfile==0.12.1 @@ -937,10 +1142,17 @@ sqlalchemy==2.0.41 # optuna sqlitedict==2.1.0 # via lm-eval -starlette==0.50.0 +sse-starlette==3.4.5 + # via mcp +starlette==1.3.1 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi - # schemathesis + # mcp + # model-hosting-container-standards + # prometheus-fastapi-instrumentator + # sse-starlette # starlette-testclient starlette-testclient==0.4.1 # via schemathesis @@ -948,6 +1160,8 @@ statsmodels==0.14.4 # via genai-perf structlog==25.4.0 # via gpt-oss +supervisor==4.3.0 + # via model-hosting-container-standards sympy==1.13.3 # via # einx @@ -965,6 +1179,7 @@ tenacity==9.1.2 # gpt-oss # lm-eval # plotly + # schemathesis tensorizer==2.10.1 # via -r requirements/test/cuda.in termcolor==3.1.0 @@ -976,6 +1191,7 @@ tifffile==2025.3.30 tiktoken==0.12.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common @@ -987,18 +1203,16 @@ timm==1.0.17 tokenizers==0.22.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in # transformers -tomli==2.2.1 - # via schemathesis -tomli-w==1.2.0 - # via schemathesis torch==2.11.0+cu130 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in # accelerate # bitsandbytes + # compressed-tensors # encodec # instanttensor # mteb @@ -1012,12 +1226,17 @@ torch==2.11.0+cu130 # torchvision # vector-quantize-pytorch # vocos + # xgrammar torchaudio==2.11.0+cu130 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in # encodec # vocos +torchcodec==0.14.0+cu130 + # via + # -c requirements/cuda.txt + # -r requirements/test/cuda.in torchvision==0.26.0+cu130 # via # -c requirements/cuda.txt @@ -1027,6 +1246,7 @@ torchvision==0.26.0+cu130 # timm tqdm==4.67.3 # via + # -r requirements/test/../common.txt # datasets # evaluate # huggingface-hub @@ -1034,24 +1254,30 @@ tqdm==4.67.3 # mteb # nltk # open-clip-torch + # openai # optuna # peft # pqdm # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.13.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/cuda.in + # compressed-tensors # genai-perf # peft # sentence-transformers # transformers-stream-generator + # xgrammar transformers-stream-generator==0.0.5 # via -r requirements/test/cuda.in triton==3.6.0 - # via torch + # via + # torch + # xgrammar tritonclient==2.64.0 # via -r requirements/test/cuda.in typepy==1.3.2 @@ -1059,20 +1285,23 @@ typepy==1.3.2 # dataproperty # pytablewriter # tabledata -typer==0.15.2 +typer==0.26.8 # via + # fastapi-cli + # fastapi-cloud-cli # fastsafetensors - # huggingface-hub # perceptron # transformers -types-python-dateutil==2.9.0.20241206 - # via arrow typing-extensions==4.15.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiosignal # albumentations # alembic + # anthropic + # anyio + # apache-tvm-ffi # azure-core # azure-identity # azure-storage-blob @@ -1082,27 +1311,36 @@ typing-extensions==4.15.0 # huggingface-hub # librosa # lm-eval + # mcp # mistral-common # mteb + # openai # opentelemetry-api + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http # opentelemetry-sdk # opentelemetry-semantic-conventions # pqdm # pydantic # pydantic-core # pydantic-extra-types + # pytest-asyncio + # rich-toolkit + # schemathesis # sentence-transformers # sqlalchemy # starlette # torch - # typer # typing-inspection + # xgrammar typing-inspection==0.4.2 - # via pydantic + # via + # fastapi + # mcp + # pydantic + # pydantic-settings tzdata==2024.2 # via pandas -uri-template==1.3.0 - # via jsonschema urllib3==2.2.3 # via # blobfile @@ -1110,32 +1348,46 @@ urllib3==2.2.3 # docker # requests # responses + # sentry-sdk # tritonclient uvicorn==0.35.0 - # via gpt-oss + # via + # fastapi + # fastapi-cli + # fastapi-cloud-cli + # gpt-oss + # mcp +uvloop==0.22.1 + # via uvicorn vector-quantize-pytorch==1.21.2 # via -r requirements/test/cuda.in virtualenv==20.31.2 # via ray vocos==0.1.0 # via -r requirements/test/cuda.in +watchfiles==1.2.0 + # via + # -r requirements/test/../common.txt + # uvicorn wcwidth==0.2.13 # via ftfy -webcolors==24.11.1 - # via jsonschema +websockets==16.0 + # via uvicorn werkzeug==3.1.3 # via schemathesis word2number==1.1 # via lm-eval wrapt==1.17.2 # via smart-open +xgrammar==0.2.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt xxhash==3.5.0 # via # datasets # evaluate yarl==1.17.1 - # via - # aiohttp - # schemathesis + # via aiohttp zipp==3.23.0 # via importlib-metadata diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 10eb7a62191..dfefa8239c5 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -23,15 +23,15 @@ jiwer # required for audio tests timm # required for internvl test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.11.3 # required for voxtral test +mistral_common[image,audio] >= 1.11.5 # required for voxtral test num2words # required for smolvlm test opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.13.1 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test. +schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes>=0.49.2 buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index ed10270f565..c5f8f85f2f3 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -15,12 +15,11 @@ albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests backoff # required for phi4mm test blobfile # required for kimi-vl test -einops # required for MPT, qwen-vl httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.15.0 # required for phi-4-mm test +peft>=0.19.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -30,16 +29,15 @@ tblib # for pickling test exceptions timm>=1.0.17 # required for internvl and gemma3n-mm test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio]>=1.11.3 # required for voxtral test +mistral_common[image,audio]>=1.11.5 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py -opencv-python-headless>=4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.13.1 tokenizers==0.22.2 -schemathesis>=3.39.15 # Required for openai schema test +schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 @@ -54,11 +52,9 @@ grpcio-reflection==1.78.0 arctic-inference==0.1.1 # Required for suffix decoding test numba==0.65.0 # Required for N-gram speculative decoding -numpy runai-model-streamer[s3,gcs,azure]==0.15.7 fastsafetensors>=0.3.2 instanttensor>=0.1.5 -pydantic>=2.12 # 2.11 leads to error on python 3.13 decord==0.6.0 # Prithvi tests @@ -74,6 +70,7 @@ gpt-oss>=0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank>=1.18.7 # required for fireredasr2 test +cohere_melody>=0.9.0 # required for cohere command reasoning parser test # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. # Older versions are in conflict with terratorch requirements. diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index eced7117116..1b4484c3a2d 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -44,21 +44,19 @@ anyio==4.13.0 # watchfiles apache-tvm-ffi==0.1.10 # via + # -c requirements/rocm.txt # tilelang # xgrammar arctic-inference==0.1.1 # via -r requirements/test/rocm.in argcomplete==3.6.3 # via datamodel-code-generator -arrow==1.4.0 - # via isoduration astor==0.8.1 # via depyf attrs==26.1.0 # via # aiohttp # jsonschema - # pytest-subtests # referencing audioread==3.0.1 # via librosa @@ -73,9 +71,7 @@ azure-identity==1.25.3 azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 - # via - # -r requirements/test/rocm.in - # schemathesis + # via -r requirements/test/rocm.in bitsandbytes==0.49.2 # via -r requirements/test/rocm.in black==26.3.1 @@ -120,9 +116,10 @@ choreographer==1.2.1 # via kaleido chz==0.4.0 # via gpt-oss -click==8.3.1 +click==8.4.2 # via # black + # huggingface-hub # jiwer # nltk # ray @@ -134,11 +131,12 @@ cloudpickle==3.1.2 # via # -r requirements/test/../common.txt # tilelang +cohere-melody==0.9.0 + # via -r requirements/test/rocm.in colorama==0.4.6 # via # perceptron # sacrebleu - # schemathesis colorful==0.5.8 # via ray colorlog==6.10.1 @@ -210,7 +208,6 @@ docstring-parser==0.17.0 einops==0.8.2 # via # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # encodec # vector-quantize-pytorch # vocos @@ -257,8 +254,6 @@ filelock==3.25.2 # virtualenv fonttools==4.62.1 # via matplotlib -fqdn==1.5.1 - # via jsonschema frozendict==2.4.7 # via einx frozenlist==1.8.0 @@ -278,10 +273,6 @@ genai-perf==0.0.16 # via -r requirements/test/rocm.in genson==1.3.0 # via datamodel-code-generator -gguf==0.18.0 - # via - # -c requirements/common.txt - # -r requirements/test/../common.txt google-api-core==2.30.0 # via # google-cloud-core @@ -331,9 +322,9 @@ h11==0.16.0 # uvicorn h2==4.3.0 # via httpx -harfile==0.4.0 +harfile==0.5.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub hiredis==3.3.1 # via tensorizer @@ -359,7 +350,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.10.2 +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -381,7 +372,7 @@ hypothesis==6.151.9 # hypothesis-graphql # hypothesis-jsonschema # schemathesis -hypothesis-graphql==0.12.0 +hypothesis-graphql==0.13.0 # via schemathesis hypothesis-jsonschema==0.23.1 # via schemathesis @@ -390,7 +381,6 @@ idna==3.11 # anyio # email-validator # httpx - # jsonschema # requests # yarl ijson==3.5.0 @@ -411,8 +401,6 @@ interegular==0.3.3 # via lm-format-enforcer isodate==0.7.2 # via azure-storage-blob -isoduration==20.11.0 - # via jsonschema isort==8.0.1 # via datamodel-code-generator jinja2==3.1.6 @@ -438,15 +426,16 @@ joblib==1.5.3 # librosa # nltk # scikit-learn -jsonpointer==3.1.0 - # via jsonschema jsonschema==4.26.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # hypothesis-jsonschema # mcp # mistral-common # ray - # schemathesis +jsonschema-rs==0.46.5 + # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema junit-xml==1.9 @@ -512,7 +501,7 @@ mcp==1.27.0 # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -574,7 +563,6 @@ numba==0.65.0 numpy==2.2.6 # via # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # accelerate # albumentations # bitsandbytes @@ -588,7 +576,6 @@ numpy==2.2.6 # evaluate # fastparquet # genai-perf - # gguf # imagehash # imageio # librosa @@ -644,7 +631,6 @@ opencv-python-headless==4.13.0.92 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # albumentations # mistral-common openpyxl==3.1.5 @@ -745,7 +731,7 @@ pathvalidate==3.3.1 # via pytablewriter patsy==1.0.2 # via statsmodels -peft==0.18.1 +peft==0.19.1 # via -r requirements/test/rocm.in perceptron==0.1.4 # via -r requirements/test/rocm.in @@ -794,7 +780,7 @@ prometheus-client==0.24.1 # opentelemetry-exporter-prometheus # prometheus-fastapi-instrumentator # ray -prometheus-fastapi-instrumentator==7.1.0 +prometheus-fastapi-instrumentator==8.0.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -848,7 +834,6 @@ pydantic==2.12.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # -r requirements/test/rocm.in # albumentations # anthropic # compressed-tensors @@ -878,20 +863,22 @@ pydantic-settings==2.13.1 # fastapi # mcp pygments==2.19.2 - # via rich + # via + # pytest + # rich pyjwt==2.12.1 # via # mcp # msal pyparsing==3.3.2 # via matplotlib -pyrate-limiter==3.9.0 +pyrate-limiter==4.4.0 # via schemathesis pystemmer==3.0.0 # via mteb pytablewriter==1.2.1 # via lm-eval -pytest==8.3.5 +pytest==9.1.0 # via # -r requirements/test/rocm.in # buildkite-test-collector @@ -902,10 +889,9 @@ pytest==8.3.5 # pytest-mock # pytest-rerunfailures # pytest-shard - # pytest-subtests # pytest-timeout # schemathesis -pytest-asyncio==0.24.0 +pytest-asyncio==1.4.0 # via -r requirements/test/rocm.in pytest-cov==6.3.0 # via -r requirements/test/rocm.in @@ -917,13 +903,10 @@ pytest-rerunfailures==14.0 # via -r requirements/test/rocm.in pytest-shard==0.1.2 # via -r requirements/test/rocm.in -pytest-subtests==0.14.2 - # via schemathesis pytest-timeout==2.3.1 # via -r requirements/test/rocm.in python-dateutil==2.9.0.post0 # via - # arrow # botocore # matplotlib # pandas @@ -958,7 +941,6 @@ pyyaml==6.0.3 # datamodel-code-generator # datasets # genai-perf - # gguf # huggingface-hub # lm-format-enforcer # optuna @@ -1003,7 +985,6 @@ requests==2.32.5 # datasets # docker # evaluate - # gguf # google-api-core # google-cloud-storage # gpt-oss @@ -1020,16 +1001,13 @@ requests==2.32.5 # tiktoken responses==0.26.0 # via genai-perf -rfc3339-validator==0.1.4 - # via jsonschema -rfc3987==1.3.8 - # via jsonschema rich==14.3.3 # via # genai-perf # mteb # perceptron # rich-toolkit + # schemathesis # typer rich-toolkit==0.19.7 # via @@ -1057,7 +1035,7 @@ s3transfer==0.16.0 # via boto3 sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1067,7 +1045,7 @@ safetensors==0.7.0 # segmentation-models-pytorch # timm # transformers -schemathesis==3.39.15 +schemathesis==4.21.6 # via -r requirements/test/rocm.in scikit-image==0.26.0 # via albumentations @@ -1124,7 +1102,6 @@ six==1.17.0 # junit-xml # opencensus # python-dateutil - # rfc3339-validator # rouge-score smart-open==7.5.1 # via ray @@ -1153,13 +1130,14 @@ sqlitedict==2.1.0 # via lm-eval sse-starlette==3.3.4 # via mcp -starlette==0.52.1 +starlette==1.3.1 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi # mcp # model-hosting-container-standards # prometheus-fastapi-instrumentator - # schemathesis # sse-starlette # starlette-testclient starlette-testclient==0.4.1 @@ -1186,6 +1164,7 @@ tenacity==9.1.4 # via # gpt-oss # lm-eval + # schemathesis tensorizer==2.10.1 # via # -c requirements/rocm.txt @@ -1219,10 +1198,6 @@ tokenizers==0.22.2 # -r requirements/test/../common.txt # -r requirements/test/rocm.in # transformers -tomli==2.4.0 - # via schemathesis -tomli-w==1.2.0 - # via schemathesis torch-c-dlpack-ext==0.1.5 # via tilelang tqdm==4.67.3 @@ -1230,7 +1205,6 @@ tqdm==4.67.3 # -r requirements/test/../common.txt # datasets # evaluate - # gguf # huggingface-hub # lm-eval # mteb @@ -1244,7 +1218,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.5.3 +transformers==5.13.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1271,7 +1245,6 @@ typer==0.24.1 # fastapi-cli # fastapi-cloud-cli # fastsafetensors - # huggingface-hub # perceptron # transformers typing-extensions==4.15.0 @@ -1306,8 +1279,10 @@ typing-extensions==4.15.0 # pydantic # pydantic-core # pydantic-extra-types + # pytest-asyncio # referencing # rich-toolkit + # schemathesis # sentence-transformers # sqlalchemy # starlette @@ -1322,10 +1297,6 @@ typing-inspection==0.4.2 # mcp # pydantic # pydantic-settings -tzdata==2025.3 - # via arrow -uri-template==1.3.0 - # via jsonschema urllib3==2.6.3 # via # blobfile @@ -1356,8 +1327,6 @@ watchfiles==1.1.1 # uvicorn wcwidth==0.6.0 # via ftfy -webcolors==25.10.0 - # via jsonschema websockets==16.0 # via uvicorn werkzeug==3.1.6 @@ -1366,7 +1335,7 @@ word2number==1.1 # via lm-eval wrapt==2.1.2 # via smart-open -xgrammar==0.2.0 +xgrammar==0.2.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1375,9 +1344,7 @@ xxhash==3.6.0 # datasets # evaluate yarl==1.23.0 - # via - # aiohttp - # schemathesis + # via aiohttp z3-solver==4.15.4.0 # via tilelang zipp==3.23.0 diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index a828867845f..d2380a75131 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -1,3 +1,5 @@ +-r ../common.txt + # --- Test Infrastructure --- tblib pytest @@ -13,8 +15,9 @@ pytest-shard absl-py accelerate arctic-inference -lm_eval[api] -modelscope +lm_eval[api]>=0.4.12 +modelscope<1.38 +transformers==5.13.1 # --- Audio Processing --- librosa diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6d5435462ff..cd745684119 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -11,6 +11,7 @@ aiohappyeyeballs==2.6.1 aiohttp==3.13.4 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # fsspec # gpt-oss # lm-eval @@ -24,22 +25,36 @@ annotated-doc==0.0.4 # typer annotated-types==0.7.0 # via pydantic +anthropic==0.112.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt anyio==4.13.0 # via + # anthropic # httpx + # mcp + # openai + # sse-starlette # starlette + # watchfiles +apache-tvm-ffi==0.1.12 + # via xgrammar arctic-inference==0.1.1 # via -r requirements/test/xpu.in +astor==0.8.1 + # via depyf attrs==26.1.0 # via # aiohttp - # jsonlines # jsonschema # referencing audioread==3.0.1 # via # -r requirements/test/xpu.in # librosa +blake3==1.0.9 + # via -r requirements/test/../common.txt blobfile==3.0.0 # via -r requirements/test/xpu.in bm25s==0.2.13 @@ -48,30 +63,47 @@ bm25s==0.2.13 # mteb bounded-pool-executor==0.0.3 # via pqdm +cachetools==7.1.4 + # via -r requirements/test/../common.txt +cbor2==6.1.2 + # via -r requirements/test/../common.txt certifi==2026.2.25 # via # httpcore # httpx # requests + # sentry-sdk cffi==2.0.0 - # via soundfile + # via + # cryptography + # soundfile chardet==5.2.0 # via mbstrdecoder charset-normalizer==3.4.6 # via requests chz==0.4.0 # via gpt-oss -click==8.3.1 +click==8.4.2 # via + # huggingface-hub # jiwer # nltk + # rich-toolkit # schemathesis # typer # uvicorn +cloudpickle==3.1.2 + # via -r requirements/test/../common.txt colorama==0.4.6 # via sacrebleu +compressed-tensors==0.17.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt coverage==7.13.5 # via pytest-cov +cryptography==49.0.0 + # via pyjwt dataproperty==1.1.0 # via # pytablewriter @@ -83,16 +115,35 @@ datasets==4.8.4 # mteb decorator==5.2.1 # via librosa +depyf==0.20.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +detect-installer==0.1.0 + # via fastapi-cloud-cli dill==0.4.1 # via # datasets + # depyf # evaluate # lm-eval # multiprocess +diskcache==5.6.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +distro==1.9.0 + # via + # anthropic + # openai +dnspython==2.8.0 + # via email-validator docker==7.1.0 # via gpt-oss docopt==0.6.2 # via num2words +docstring-parser==0.18.0 + # via anthropic dpcpp-cpp-rt==2025.3.2 # via # onemkl-sycl-blas @@ -101,15 +152,30 @@ dpcpp-cpp-rt==2025.3.2 # onemkl-sycl-rng # onemkl-sycl-sparse # torch +einops==0.8.2 + # via -r requirements/test/../common.txt +email-validator==2.3.0 + # via + # fastapi + # pydantic evaluate==0.4.6 # via lm-eval fastapi==0.135.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss + # model-hosting-container-standards +fastapi-cli==0.0.27 + # via fastapi +fastapi-cloud-cli==0.21.0 + # via fastapi-cli +fastar==0.11.0 + # via fastapi-cloud-cli filelock==3.25.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # blobfile # datasets # huggingface-hub @@ -125,28 +191,44 @@ fsspec==2026.2.0 # evaluate # huggingface-hub # torch +googleapis-common-protos==1.75.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http gpt-oss==0.0.8 # via -r requirements/test/xpu.in graphql-core==3.2.8 # via hypothesis-graphql +grpcio==1.81.1 + # via opentelemetry-exporter-otlp-proto-grpc h11==0.16.0 # via # httpcore # uvicorn harfile==0.4.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub html2text==2025.4.15 # via gpt-oss httpcore==1.0.9 # via httpx +httptools==0.8.0 + # via uvicorn httpx==0.28.1 # via + # anthropic # datasets + # fastapi + # fastapi-cloud-cli # huggingface-hub + # mcp + # model-hosting-container-standards + # openai # schemathesis -huggingface-hub==1.10.2 +httpx-sse==0.4.3 + # via mcp +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -167,9 +249,12 @@ hypothesis-jsonschema==0.23.1 idna==3.11 # via # anyio + # email-validator # httpx # requests # yarl +ijson==3.5.0 + # via -r requirements/test/../common.txt imageio==2.37.3 # via scikit-image impi-rt==2021.17.2 @@ -213,23 +298,33 @@ intel-sycl-rt==2025.3.2 # dpcpp-cpp-rt # oneccl # torch +interegular==0.3.3 + # via lm-format-enforcer jinja2==3.1.6 # via # -c requirements/xpu.txt + # fastapi # lm-eval # torch +jiter==0.15.0 + # via + # anthropic + # openai jiwer==4.0.0 # via -r requirements/test/xpu.in +jmespath==1.1.0 + # via model-hosting-container-standards joblib==1.5.3 # via # librosa # nltk # scikit-learn -jsonlines==4.0.0 - # via lm-eval jsonschema==4.26.0 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # hypothesis-jsonschema + # mcp # mistral-common # schemathesis jsonschema-rs==0.45.0 @@ -238,16 +333,30 @@ jsonschema-specifications==2025.9.1 # via jsonschema junit-xml==1.9 # via schemathesis +lark==1.2.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt lazy-loader==0.5 # via # librosa # scikit-image librosa==0.10.2.post1 # via -r requirements/test/xpu.in +llguidance==1.7.6 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt llvmlite==0.47.0 # via numba -lm-eval==0.4.11 +lm-eval==0.4.12 # via -r requirements/test/xpu.in +lm-format-enforcer==0.11.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +loguru==0.7.3 + # via compressed-tensors lxml==6.0.2 # via # blobfile @@ -264,11 +373,14 @@ mbstrdecoder==1.1.4 # dataproperty # pytablewriter # typepy +mcp==1.28.1 + # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.11.3 +mistral-common==1.11.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # -r requirements/test/xpu.in mkl==2025.3.1 # via @@ -278,6 +390,10 @@ mkl==2025.3.1 # onemkl-sycl-rng # onemkl-sycl-sparse # torch +model-hosting-container-standards==0.1.16 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt modelscope==1.35.3 # via -r requirements/test/xpu.in more-itertools==10.8.0 @@ -286,6 +402,8 @@ mpmath==1.3.0 # via sympy msgpack==1.1.2 # via librosa +msgspec==0.21.1 + # via -r requirements/test/../common.txt mteb==2.12.7 # via -r requirements/test/xpu.in multidict==6.7.1 @@ -300,6 +418,8 @@ networkx==3.6.1 # via # scikit-image # torch +ninja==1.13.0 + # via -r requirements/test/../common.txt nltk==3.9.4 # via rouge-score num2words==0.5.14 @@ -310,6 +430,7 @@ numba==0.65.0 # librosa numpy==2.2.6 # via + # -r requirements/test/../common.txt # accelerate # albumentations # bm25s @@ -335,6 +456,7 @@ numpy==2.2.6 # tifffile # torchvision # transformers + # xgrammar oneccl==2021.17.2 # via # oneccl-devel @@ -358,15 +480,65 @@ onemkl-sycl-rng==2025.3.1 # via torch onemkl-sycl-sparse==2025.3.1 # via torch +openai==2.44.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt openai-harmony==0.0.8 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss opencv-python-headless==4.13.0.92 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # albumentations # mistral-common +opentelemetry-api==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions +opentelemetry-exporter-otlp==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +opentelemetry-exporter-otlp-proto-common==1.43.0 + # via + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-exporter-otlp-proto-grpc==1.43.0 + # via opentelemetry-exporter-otlp +opentelemetry-exporter-otlp-proto-http==1.43.0 + # via opentelemetry-exporter-otlp +opentelemetry-proto==1.43.0 + # via + # opentelemetry-exporter-otlp-proto-common + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http +opentelemetry-sdk==1.43.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-semantic-conventions-ai +opentelemetry-semantic-conventions==0.64b0 + # via + # opentelemetry-sdk + # opentelemetry-semantic-conventions-ai +opentelemetry-semantic-conventions-ai==0.5.1 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt +outlines-core==0.2.14 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt packaging==26.0 # via # -c requirements/xpu.txt @@ -375,6 +547,7 @@ packaging==26.0 # evaluate # huggingface-hub # lazy-loader + # lm-format-enforcer # modelscope # pooch # pytest @@ -386,10 +559,13 @@ pandas==3.0.1 # via # datasets # evaluate +partial-json-parser==0.2.1.1.post7 + # via -r requirements/test/../common.txt pathvalidate==3.3.1 # via pytablewriter pillow==12.1.1 # via + # -r requirements/test/../common.txt # imageio # mistral-common # scikit-image @@ -412,16 +588,37 @@ portalocker==3.2.0 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/xpu.in +prometheus-client==0.25.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # prometheus-fastapi-instrumentator +prometheus-fastapi-instrumentator==8.0.2 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt propcache==0.4.1 # via # aiohttp # yarl +protobuf==7.35.1 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt + # googleapis-common-protos + # opentelemetry-proto psutil==7.2.2 - # via accelerate + # via + # -r requirements/test/../common.txt + # accelerate py==1.11.0 # via pytest-forked +py-cpuinfo==9.0.0 + # via -r requirements/test/../common.txt pyarrow==23.0.1 # via datasets +pybase64==1.4.3 + # via -r requirements/test/../common.txt pycountry==26.2.16 # via pydantic-extra-types pycparser==3.0 @@ -431,23 +628,41 @@ pycryptodomex==3.23.0 pydantic==2.12.5 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # albumentations + # anthropic + # compressed-tensors # fastapi + # fastapi-cloud-cli # gpt-oss + # lm-format-enforcer + # mcp # mistral-common + # model-hosting-container-standards # mteb + # openai # openai-harmony # pydantic-extra-types + # pydantic-settings + # xgrammar pydantic-core==2.41.5 # via pydantic pydantic-extra-types==2.11.1 - # via mistral-common + # via + # fastapi + # mistral-common +pydantic-settings==2.14.2 + # via + # fastapi + # mcp pyelftools==0.32 # via triton-xpu pygments==2.20.0 # via # pytest # rich +pyjwt==2.13.0 + # via mcp pyrate-limiter==4.1.0 # via schemathesis pystemmer==3.0.0 @@ -482,19 +697,36 @@ python-dateutil==2.9.0.post0 # via # pandas # typepy +python-dotenv==1.2.2 + # via + # pydantic-settings + # uvicorn +python-json-logger==4.1.0 + # via -r requirements/test/../common.txt +python-multipart==0.0.32 + # via + # fastapi + # mcp pytrec-eval-terrier==0.5.10 # via mteb pytz==2026.1.post1 # via typepy pyyaml==6.0.3 # via + # -r requirements/test/../common.txt # accelerate # albumentations # datasets # huggingface-hub + # lm-format-enforcer # schemathesis # timm # transformers + # uvicorn +pyzmq==27.1.0 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt rapidfuzz==3.12.1 # via # -r requirements/test/xpu.in @@ -505,6 +737,7 @@ referencing==0.37.0 # jsonschema-specifications regex==2026.3.32 # via + # -r requirements/test/../common.txt # nltk # sacrebleu # tiktoken @@ -512,6 +745,7 @@ regex==2026.3.32 requests==2.33.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # datasets # docker # evaluate @@ -520,6 +754,7 @@ requests==2.33.1 # mistral-common # modelscope # mteb + # opentelemetry-exporter-otlp-proto-http # pooch # schemathesis # starlette-testclient @@ -527,8 +762,15 @@ requests==2.33.1 rich==14.3.3 # via # mteb + # rich-toolkit # schemathesis # typer +rich-toolkit==0.20.1 + # via + # fastapi-cli + # fastapi-cloud-cli +rignore==0.7.6 + # via fastapi-cloud-cli rouge-score==0.1.2 # via lm-eval rpds-py==0.30.0 @@ -537,9 +779,10 @@ rpds-py==0.30.0 # referencing sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # accelerate # timm # transformers @@ -566,10 +809,18 @@ scipy==1.17.1 # sentence-transformers sentence-transformers==5.3.0 # via mteb +sentencepiece==0.2.1 + # via -r requirements/test/../common.txt +sentry-sdk==2.63.0 + # via fastapi-cloud-cli +setproctitle==1.3.7 + # via -r requirements/test/../common.txt setuptools==80.10.2 # via # -c requirements/common.txt # -c requirements/xpu.txt + # -r requirements/test/../common.txt + # model-hosting-container-standards # modelscope # pytablewriter # torch @@ -578,9 +829,14 @@ shellingham==1.5.4 six==1.17.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # junit-xml # python-dateutil # rouge-score +sniffio==1.3.1 + # via + # anthropic + # openai sortedcontainers==2.4.0 # via hypothesis soundfile==0.13.1 @@ -595,14 +851,24 @@ soxr==0.5.0.post1 # mistral-common sqlitedict==2.1.0 # via lm-eval -starlette==1.0.0 +sse-starlette==3.4.5 + # via mcp +starlette==1.3.1 # via + # -c requirements/common.txt + # -r requirements/test/../common.txt # fastapi + # mcp + # model-hosting-container-standards + # prometheus-fastapi-instrumentator + # sse-starlette # starlette-testclient starlette-testclient==0.4.1 # via schemathesis structlog==25.5.0 # via gpt-oss +supervisor==4.3.0 + # via model-hosting-container-standards sympy==1.14.0 # via torch tabledata==1.3.4 @@ -637,6 +903,7 @@ tifffile==2026.3.3 tiktoken==0.12.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common @@ -645,19 +912,23 @@ timm==1.0.17 tokenizers==0.22.2 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # transformers torch==2.12.0+xpu # via # -c requirements/xpu.txt # accelerate + # compressed-tensors # mteb # sentence-transformers # timm # torchvision + # xgrammar torchvision==0.27.0+xpu # via timm tqdm==4.67.3 # via + # -r requirements/test/../common.txt # datasets # evaluate # huggingface-hub @@ -665,13 +936,20 @@ tqdm==4.67.3 # modelscope # mteb # nltk + # openai # pqdm # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.13.1 # via # -c requirements/common.txt + # -r requirements/test/../common.txt + # -r requirements/test/xpu.in + # compressed-tensors # sentence-transformers + # xgrammar +triton==3.7.1 + # via xgrammar triton-xpu==3.7.1 # via torch typepy==1.3.4 @@ -681,36 +959,52 @@ typepy==1.3.4 # tabledata typer==0.24.1 # via - # huggingface-hub + # fastapi-cli + # fastapi-cloud-cli # transformers typing-extensions==4.15.0 # via # -c requirements/common.txt + # -r requirements/test/../common.txt # aiosignal # albumentations + # anthropic # anyio + # apache-tvm-ffi # chz # fastapi + # grpcio # huggingface-hub # librosa # lm-eval + # mcp # mistral-common # mteb + # openai + # opentelemetry-api + # opentelemetry-exporter-otlp-proto-grpc + # opentelemetry-exporter-otlp-proto-http + # opentelemetry-sdk + # opentelemetry-semantic-conventions # pqdm # pydantic # pydantic-core # pydantic-extra-types # pytest-asyncio # referencing + # rich-toolkit # schemathesis # sentence-transformers # starlette # torch # typing-inspection + # xgrammar typing-inspection==0.4.2 # via # fastapi + # mcp # pydantic + # pydantic-settings umf==1.0.3 # via # intel-cmplr-lib-ur @@ -721,17 +1015,33 @@ urllib3==2.6.3 # docker # modelscope # requests + # sentry-sdk uvicorn==0.42.0 - # via gpt-oss + # via + # fastapi + # fastapi-cli + # fastapi-cloud-cli + # gpt-oss + # mcp +uvloop==0.22.1 + # via uvicorn +watchfiles==1.2.0 + # via + # -r requirements/test/../common.txt + # uvicorn +websockets==16.0 + # via uvicorn werkzeug==3.1.7 # via schemathesis word2number==1.1 # via lm-eval +xgrammar==0.2.3 + # via + # -c requirements/common.txt + # -r requirements/test/../common.txt xxhash==3.6.0 # via # datasets # evaluate yarl==1.23.0 # via aiohttp -zstandard==0.25.0 - # via lm-eval diff --git a/requirements/tpu.txt b/requirements/tpu.txt index ea2c24bba0f..da477a68461 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.21.0 +tpu-inference==0.24.0 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 33a808866c4..2da3c6654b1 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -15,6 +15,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding torch==2.12.0 torchaudio torchvision +torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib>=0.13.0 -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl +auto_round_lib>=0.14.0 +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ef8c2b90a15..d155e07edad 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -31,24 +31,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "aligned" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685" -dependencies = [ - "as-slice", -] - -[[package]] -name = "aligned-vec" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" -dependencies = [ - "equator", -] - [[package]] name = "android_system_properties" version = "0.1.5" @@ -144,12 +126,6 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" -[[package]] -name = "arbitrary" -version = "1.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" - [[package]] name = "arc-swap" version = "1.9.0" @@ -159,17 +135,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "arg_enum_proc_macro" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "arrayref" version = "0.3.9" @@ -182,15 +147,6 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" -[[package]] -name = "as-slice" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516" -dependencies = [ - "stable_deref_trait", -] - [[package]] name = "async-io" version = "2.6.0" @@ -316,55 +272,35 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "auto_enums" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e4487600931c9a89f8db7ffbdf3fbdd45bb7bd85e26861f659a463cd0dff966" +dependencies = [ + "derive_utils", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "auto_impl" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdcb70bdbc4d478427380519163274ac86e52916e10f0a8889adf0f96d3fee7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "autocfg" version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "av-scenechange" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f321d77c20e19b92c39e7471cf986812cbb46659d2af674adc4331ef3f18394" -dependencies = [ - "aligned", - "anyhow", - "arg_enum_proc_macro", - "arrayvec", - "log", - "num-rational", - "num-traits", - "pastey", - "rayon", - "thiserror 2.0.18", - "v_frame", - "y4m", -] - -[[package]] -name = "av1-grain" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cfddb07216410377231960af4fcab838eaa12e013417781b78bd95ee22077f8" -dependencies = [ - "anyhow", - "arrayvec", - "log", - "nom 8.0.0", - "num-rational", - "v_frame", -] - -[[package]] -name = "avif-serialize" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "375082f007bd67184fb9c0374614b29f9aaa604ec301635f72338bb65386a53d" -dependencies = [ - "arrayvec", -] - [[package]] name = "axum" version = "0.8.8" @@ -443,12 +379,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "base64ct" -version = "1.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" - [[package]] name = "bit-set" version = "0.5.3" @@ -479,27 +409,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" -[[package]] -name = "bit_field" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6" - [[package]] name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" -[[package]] -name = "bitstream-io" -version = "4.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eff00be299a18769011411c9def0d827e8f2d7bf0c3dbf53633147a8867fd1f" -dependencies = [ - "no_std_io2", -] - [[package]] name = "blake3" version = "1.8.5" @@ -534,12 +449,6 @@ dependencies = [ "serde", ] -[[package]] -name = "built" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4ad8f11f288f48ca24471bbd51ac257aaeaaa07adae295591266b792902ae64" - [[package]] name = "bumpalo" version = "3.20.2" @@ -580,9 +489,9 @@ checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" dependencies = [ "serde", ] @@ -620,12 +529,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "cfg_aliases" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" - [[package]] name = "chrono" version = "0.4.44" @@ -743,19 +646,6 @@ dependencies = [ "crossbeam-utils", ] -[[package]] -name = "console" -version = "0.15.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" -dependencies = [ - "encode_unicode", - "libc", - "once_cell", - "unicode-width", - "windows-sys 0.59.0", -] - [[package]] name = "console" version = "0.16.2" @@ -775,35 +665,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" -[[package]] -name = "cookie" -version = "0.18.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" -dependencies = [ - "percent-encoding", - "time", - "version_check", -] - -[[package]] -name = "cookie_store" -version = "0.22.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" -dependencies = [ - "cookie", - "document-features", - "idna", - "indexmap 2.13.0", - "log", - "serde", - "serde_derive", - "serde_json", - "time", - "url", -] - [[package]] name = "core-foundation" version = "0.9.4" @@ -1027,16 +888,6 @@ dependencies = [ "serde", ] -[[package]] -name = "der" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b" -dependencies = [ - "pem-rfc7468", - "zeroize", -] - [[package]] name = "deranged" version = "0.5.8" @@ -1099,6 +950,17 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "derive_utils" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "362f47930db19fe7735f527e6595e4900316b893ebf6d48ad3d31be928d57dd6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "digest" version = "0.10.7" @@ -1263,26 +1125,6 @@ dependencies = [ "log", ] -[[package]] -name = "equator" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" -dependencies = [ - "equator-macro", -] - -[[package]] -name = "equator-macro" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "equivalent" version = "1.0.2" @@ -1315,7 +1157,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" dependencies = [ "futures-core", - "nom 7.1.3", + "nom", "pin-project-lite", ] @@ -1329,21 +1171,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "exr" -version = "1.74.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4300e043a56aa2cb633c01af81ca8f699a321879a7854d3896a0ba89056363be" -dependencies = [ - "bit_field", - "half", - "lebe", - "miniz_oxide", - "rayon-core", - "smallvec", - "zune-inflate", -] - [[package]] name = "fancy-regex" version = "0.13.0" @@ -1382,13 +1209,12 @@ dependencies = [ [[package]] name = "fastokens" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "796a262ed47d1458a4b40d0ed831c927e6f54d5b9c1de2683bb4ac9b04f4c7cc" +checksum = "8728655e193e0d08d7a95d63cf1fdb9b768d282cab0a112ecb006615bae9f067" dependencies = [ "daachorse", "fancy-regex 0.17.0", - "hf-hub 0.4.3", "icu_normalizer", "memchr", "pcre2", @@ -1632,10 +1458,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", - "js-sys", "libc", "wasi", - "wasm-bindgen", ] [[package]] @@ -1677,9 +1501,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" dependencies = [ "atomic-waker", "bytes", @@ -1754,26 +1578,6 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" -[[package]] -name = "hf-hub" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" -dependencies = [ - "dirs", - "http", - "indicatif 0.17.11", - "libc", - "log", - "rand 0.9.2", - "reqwest", - "serde", - "serde_json", - "thiserror 2.0.18", - "ureq 2.12.1", - "windows-sys 0.60.2", -] - [[package]] name = "hf-hub" version = "0.5.0" @@ -1782,11 +1586,9 @@ checksum = "aef3982638978efa195ff11b305f51f1f22f4f0a6cabee7af79b383ebee6a213" dependencies = [ "dirs", "futures", - "http", - "indicatif 0.18.4", + "indicatif", "libc", "log", - "native-tls", "num_cpus", "rand 0.9.2", "reqwest", @@ -1794,7 +1596,6 @@ dependencies = [ "serde_json", "thiserror 2.0.18", "tokio", - "ureq 3.3.0", "windows-sys 0.61.2", ] @@ -1860,9 +1661,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "1.8.1" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" dependencies = [ "atomic-waker", "bytes", @@ -1875,7 +1676,6 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "pin-utils", "smallvec", "tokio", "want", @@ -1891,12 +1691,10 @@ dependencies = [ "hyper", "hyper-util", "rustls", - "rustls-native-certs", "rustls-pki-types", "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.6", ] [[package]] @@ -2103,16 +1901,11 @@ dependencies = [ "bytemuck", "byteorder-lite", "color_quant", - "exr", "gif", "image-webp", "moxcms", "num-traits", "png", - "qoi", - "ravif", - "rayon", - "rgb", "tiff", "zune-core", "zune-jpeg", @@ -2128,12 +1921,6 @@ dependencies = [ "quick-error", ] -[[package]] -name = "imgref" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7c5cedc30da3a610cac6b4ba17597bdf7152cf974e8aab3afb3d54455e371c8" - [[package]] name = "indexmap" version = "1.9.3" @@ -2157,26 +1944,13 @@ dependencies = [ "serde_core", ] -[[package]] -name = "indicatif" -version = "0.17.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" -dependencies = [ - "console 0.15.11", - "number_prefix", - "portable-atomic", - "unicode-width", - "web-time", -] - [[package]] name = "indicatif" version = "0.18.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ - "console 0.16.2", + "console", "portable-atomic", "unicode-width", "unit-prefix", @@ -2192,17 +1966,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "interpolate_name" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "ipnet" version = "2.12.0" @@ -2343,12 +2106,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" -[[package]] -name = "lebe" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" - [[package]] name = "libc" version = "0.2.183" @@ -2356,13 +2113,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" [[package]] -name = "libfuzzer-sys" -version = "0.4.12" +name = "libloading" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" dependencies = [ - "arbitrary", - "cc", + "cfg-if", + "windows-link", ] [[package]] @@ -2409,21 +2166,27 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" -version = "1.5.0" -source = "git+https://github.com/vllm-project/llm-multimodal?rev=5b558989844d1c7af3e43d0f604069ffd9c06320#5b558989844d1c7af3e43d0f604069ffd9c06320" +version = "1.7.1" +source = "git+https://github.com/smg-project/llm-multimodal?rev=c8a29dcc755139fdc26185f400ea48c6d6d48273#c8a29dcc755139fdc26185f400ea48c6d6d48273" dependencies = [ + "anyhow", "base64 0.22.1", "blake3", "bytes", "fast_image_resize", + "hf-hub", "image", + "libloading", "ndarray 0.17.2", "once_cell", "reqwest", "serde", "serde_json", + "serde_with", + "tempfile", "thiserror 2.0.18", "tokio", + "tracing", "url", ] @@ -2442,21 +2205,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "loop9" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fae87c125b03c1d2c0150c90365d7d6bcc53fb73a9acaef207d2d065860f062" -dependencies = [ - "imgref", -] - -[[package]] -name = "lru-slab" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" - [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -2556,16 +2304,6 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "maybe-rayon" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519" -dependencies = [ - "cfg-if", - "rayon", -] - [[package]] name = "memchr" version = "2.8.0" @@ -2643,9 +2381,9 @@ dependencies = [ [[package]] name = "mio" -version = "1.1.1" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "wasi", @@ -2735,21 +2473,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "rawpointer", -] - -[[package]] -name = "new_debug_unreachable" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" - -[[package]] -name = "no_std_io2" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b51ed7824b6e07d354605f4abb3d9d300350701299da96642ee084f5ce631550" -dependencies = [ - "memchr", + "serde", ] [[package]] @@ -2762,21 +2486,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "nom" -version = "8.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" -dependencies = [ - "memchr", -] - -[[package]] -name = "noop_proc_macro" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" - [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2786,16 +2495,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "num-bigint" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" -dependencies = [ - "num-integer", - "num-traits", -] - [[package]] name = "num-complex" version = "0.4.6" @@ -2811,17 +2510,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" -[[package]] -name = "num-derive" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "num-integer" version = "0.1.46" @@ -2831,17 +2519,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.19" @@ -2870,17 +2547,11 @@ dependencies = [ "libc", ] -[[package]] -name = "number_prefix" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" - [[package]] name = "once_cell" -version = "1.21.3" +version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "once_cell_polyfill" @@ -2916,30 +2587,6 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" -[[package]] -name = "openai-harmony" -version = "0.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e77e82af451fc95deeb728a40b84db8ee82d341e136c268de415123a560b9b72" -dependencies = [ - "anyhow", - "base64 0.22.1", - "bstr", - "clap", - "fancy-regex 0.13.0", - "futures", - "image", - "regex", - "reqwest", - "rustc-hash 1.1.0", - "serde", - "serde_json", - "serde_with", - "sha1", - "sha2", - "thiserror 2.0.18", -] - [[package]] name = "openai-protocol" version = "1.6.0" @@ -2961,15 +2608,14 @@ dependencies = [ [[package]] name = "openssl" -version = "0.10.76" +version = "0.10.81" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "951c002c75e16ea2c65b8c7e4d3d51d5530d8dfa7d060b4776828c88cfb18ecf" +checksum = "77823a27f0babb03091cb9ed9ef80af3b39dbc82f97e8fa530374b7dafd87a45" dependencies = [ "bitflags", "cfg-if", "foreign-types", "libc", - "once_cell", "openssl-macros", "openssl-sys", ] @@ -3002,9 +2648,9 @@ dependencies = [ [[package]] name = "openssl-sys" -version = "0.9.112" +version = "0.9.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57d55af3b3e226502be1526dfdba67ab0e9c96fc293004e79576b2b9edb0dbdb" +checksum = "b47e7e6bb2c38cd930d25a23b40fa52e068c10e85f3e03a7f5ba5aaca5713695" dependencies = [ "cc", "libc", @@ -3019,6 +2665,24 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "oss-harmony" +version = "0.0.11" +source = "git+https://github.com/oss-harmony/harmony?tag=v0.0.11#76e849426cc092f84509e31a17027755f67d662a" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.13.0", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "serde_with", + "sha2", + "thiserror 2.0.18", + "zstd", +] + [[package]] name = "parking" version = "2.2.1" @@ -3054,12 +2718,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pastey" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" - [[package]] name = "pcre2" version = "0.2.11" @@ -3082,15 +2740,6 @@ dependencies = [ "pkg-config", ] -[[package]] -name = "pem-rfc7468" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" -dependencies = [ - "base64ct", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -3172,12 +2821,6 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" -[[package]] -name = "pin-utils" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" - [[package]] name = "pkg-config" version = "0.3.32" @@ -3337,25 +2980,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "profiling" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3eb8486b569e12e2c32ad3e204dbaba5e4b5b216e9367044f25f1dba42341773" -dependencies = [ - "profiling-procmacros", -] - -[[package]] -name = "profiling-procmacros" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52717f9a02b6965224f95ca2a81e2e0c5c43baacd28ca057577988930b6c3d5b" -dependencies = [ - "quote", - "syn 2.0.117", -] - [[package]] name = "prometheus-client" version = "0.24.0" @@ -3459,12 +3083,72 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f" [[package]] -name = "qoi" -version = "0.4.1" +name = "pyo3" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f6d64c71eb498fe9eae14ce4ec935c555749aef511cca85b5568910d6e48001" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" dependencies = [ - "bytemuck", + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pythonize" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95" +dependencies = [ + "pyo3", + "serde", + "serde_json", ] [[package]] @@ -3473,61 +3157,6 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" -[[package]] -name = "quinn" -version = "0.11.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" -dependencies = [ - "bytes", - "cfg_aliases", - "pin-project-lite", - "quinn-proto", - "quinn-udp", - "rustc-hash 2.1.1", - "rustls", - "socket2", - "thiserror 2.0.18", - "tokio", - "tracing", - "web-time", -] - -[[package]] -name = "quinn-proto" -version = "0.11.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" -dependencies = [ - "bytes", - "getrandom 0.3.4", - "lru-slab", - "rand 0.9.2", - "ring", - "rustc-hash 2.1.1", - "rustls", - "rustls-pki-types", - "slab", - "thiserror 2.0.18", - "tinyvec", - "tracing", - "web-time", -] - -[[package]] -name = "quinn-udp" -version = "0.5.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" -dependencies = [ - "cfg_aliases", - "libc", - "once_cell", - "socket2", - "tracing", - "windows-sys 0.60.2", -] - [[package]] name = "quote" version = "1.0.45" @@ -3608,56 +3237,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "rav1e" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43b6dd56e85d9483277cde964fd1bdb0428de4fec5ebba7540995639a21cb32b" -dependencies = [ - "aligned-vec", - "arbitrary", - "arg_enum_proc_macro", - "arrayvec", - "av-scenechange", - "av1-grain", - "bitstream-io", - "built", - "cfg-if", - "interpolate_name", - "itertools 0.14.0", - "libc", - "libfuzzer-sys", - "log", - "maybe-rayon", - "new_debug_unreachable", - "noop_proc_macro", - "num-derive", - "num-traits", - "paste", - "profiling", - "rand 0.9.2", - "rand_chacha 0.9.0", - "simd_helpers", - "thiserror 2.0.18", - "v_frame", - "wasm-bindgen", -] - -[[package]] -name = "ravif" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e52310197d971b0f5be7fe6b57530dcd27beb35c1b013f29d66c1ad73fbbcc45" -dependencies = [ - "avif-serialize", - "imgref", - "loop9", - "quick-error", - "rav1e", - "rayon", - "rgb", -] - [[package]] name = "rawpointer" version = "0.2.1" @@ -3782,7 +3361,6 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", - "futures-channel", "futures-core", "futures-util", "h2", @@ -3800,9 +3378,6 @@ dependencies = [ "native-tls", "percent-encoding", "pin-project-lite", - "quinn", - "rustls", - "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", @@ -3810,7 +3385,6 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-native-tls", - "tokio-rustls", "tokio-util", "tower", "tower-http", @@ -3820,7 +3394,6 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 1.0.6", ] [[package]] @@ -3833,18 +3406,12 @@ dependencies = [ "futures-core", "futures-timer", "mime", - "nom 7.1.3", + "nom", "pin-project-lite", "reqwest", "thiserror 1.0.69", ] -[[package]] -name = "rgb" -version = "0.8.53" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" - [[package]] name = "ring" version = "0.17.14" @@ -3959,34 +3526,19 @@ version = "0.23.37" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" dependencies = [ - "log", "once_cell", - "ring", "rustls-pki-types", "rustls-webpki", "subtle", "zeroize", ] -[[package]] -name = "rustls-native-certs" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" -dependencies = [ - "openssl-probe", - "rustls-pki-types", - "schannel", - "security-framework", -] - [[package]] name = "rustls-pki-types" -version = "1.14.0" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" dependencies = [ - "web-time", "zeroize", ] @@ -4419,17 +3971,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "sha1" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" -dependencies = [ - "cfg-if", - "cpufeatures 0.2.17", - "digest", -] - [[package]] name = "sha2" version = "0.10.9" @@ -4472,15 +4013,6 @@ version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" -[[package]] -name = "simd_helpers" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6" -dependencies = [ - "quote", -] - [[package]] name = "siphasher" version = "1.0.2" @@ -4520,17 +4052,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "socks" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" -dependencies = [ - "byteorder", - "libc", - "winapi", -] - [[package]] name = "spm_precompiled" version = "0.1.4" @@ -4538,7 +4059,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" dependencies = [ "base64 0.13.1", - "nom 7.1.3", + "nom", "serde", "unicode-segmentation", ] @@ -4669,6 +4190,12 @@ dependencies = [ "libc", ] +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + [[package]] name = "task-local" version = "0.1.1" @@ -4891,20 +4418,21 @@ dependencies = [ ] [[package]] -name = "tinyvec" -version = "1.11.0" +name = "tls-listener" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "1461056cc1ef47003f7ee16e4cef3741068d4c7f6b627bfce49b7c00c120a530" dependencies = [ - "tinyvec_macros", + "axum", + "futures-util", + "openssl", + "pin-project-lite", + "thiserror 2.0.18", + "tokio", + "tokio-openssl", + "tracing", ] -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - [[package]] name = "tokenizers" version = "0.22.2" @@ -4918,7 +4446,7 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", - "indicatif 0.18.4", + "indicatif", "itertools 0.14.0", "log", "macro_rules_attribute", @@ -4941,9 +4469,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.50.0" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -4958,9 +4486,9 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.6.1" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", @@ -4977,6 +4505,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-openssl" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59df6849caa43bb7567f9a36f863c447d95a11d5903c9cc334ba32576a27eadd" +dependencies = [ + "openssl", + "openssl-sys", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -5443,61 +4982,6 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" -[[package]] -name = "ureq" -version = "2.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" -dependencies = [ - "base64 0.22.1", - "flate2", - "log", - "once_cell", - "rustls", - "rustls-pki-types", - "serde", - "serde_json", - "socks", - "url", - "webpki-roots 0.26.11", -] - -[[package]] -name = "ureq" -version = "3.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" -dependencies = [ - "base64 0.22.1", - "cookie_store", - "der", - "flate2", - "log", - "native-tls", - "percent-encoding", - "rustls", - "rustls-pki-types", - "serde", - "serde_json", - "socks", - "ureq-proto", - "utf8-zero", - "webpki-root-certs", - "webpki-roots 1.0.6", -] - -[[package]] -name = "ureq-proto" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" -dependencies = [ - "base64 0.22.1", - "http", - "httparse", - "log", -] - [[package]] name = "url" version = "2.5.8" @@ -5522,12 +5006,6 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" -[[package]] -name = "utf8-zero" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" - [[package]] name = "utf8_iter" version = "1.0.4" @@ -5551,17 +5029,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "v_frame" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "666b7727c8875d6ab5db9533418d7c764233ac9c0cff1d469aec8fa127597be2" -dependencies = [ - "aligned-vec", - "num-traits", - "wasm-bindgen", -] - [[package]] name = "validator" version = "0.20.0" @@ -5627,7 +5094,8 @@ dependencies = [ "llm-multimodal", "minijinja", "minijinja-contrib", - "openai-harmony", + "ndarray 0.17.2", + "oss-harmony", "paste", "reqwest", "rmp-serde", @@ -5636,10 +5104,12 @@ dependencies = [ "serde_json", "serde_with", "serial_test", + "strum", "subenum", "tempfile", "thiserror 2.0.18", "thiserror-ext", + "time", "tokio", "tracing", "tracing-subscriber", @@ -5647,10 +5117,10 @@ dependencies = [ "uuid", "vllm-engine-core-client", "vllm-llm", - "vllm-reasoning-parser", + "vllm-parser", "vllm-text", "vllm-tokenizer", - "vllm-tool-parser", + "xgrammar-structural-tag", "zeromq", ] @@ -5675,6 +5145,7 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", + "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", @@ -5729,6 +5200,7 @@ dependencies = [ "enum-as-inner", "expect-test", "futures", + "parking_lot", "rmp-serde", "serde", "serde_json", @@ -5759,6 +5231,7 @@ dependencies = [ name = "vllm-metrics" version = "0.1.0" dependencies = [ + "itertools 0.14.0", "prometheus-client", ] @@ -5782,11 +5255,22 @@ dependencies = [ ] [[package]] -name = "vllm-reasoning-parser" +name = "vllm-parser" version = "0.1.0" dependencies = [ + "criterion", + "easy-ext", + "expect-test", + "futures", + "openai-protocol", + "serde", + "serde_json", "thiserror 2.0.18", + "thiserror-ext", + "tool-parser", "vllm-tokenizer", + "winnow", + "xgrammar-structural-tag", ] [[package]] @@ -5796,6 +5280,7 @@ dependencies = [ "anyhow", "async-openai", "asynk-strim-attr", + "auto_enums", "axum", "bytes", "clap", @@ -5803,9 +5288,13 @@ dependencies = [ "expect-test", "futures", "http-body", + "hyper", + "hyper-util", + "indexmap 2.13.0", "itertools 0.14.0", "libc", "llm-multimodal", + "openssl", "prost", "prost-types", "rmp-serde", @@ -5817,8 +5306,11 @@ dependencies = [ "sha2", "socket2", "subtle", + "tempfile", "thiserror-ext", + "tls-listener", "tokio", + "tokio-openssl", "tokio-stream", "tokio-util", "tonic", @@ -5836,6 +5328,7 @@ dependencies = [ "vllm-llm", "vllm-metrics", "vllm-text", + "vllm-tokenizer", "zeromq", ] @@ -5849,8 +5342,9 @@ dependencies = [ "enum-as-inner", "expect-test", "futures", - "hf-hub 0.5.0", + "hf-hub", "itertools 0.14.0", + "reqwest", "serde", "serde_json", "serde_with", @@ -5873,7 +5367,8 @@ dependencies = [ "base64 0.22.1", "criterion", "fastokens", - "hf-hub 0.5.0", + "hf-hub", + "reqwest", "riptoken", "rustc-hash 1.1.0", "serde", @@ -5884,24 +5379,19 @@ dependencies = [ "thiserror-ext", "tiktoken-rs 0.9.1", "tokenizers", + "tokio", "tracing", ] [[package]] -name = "vllm-tool-parser" +name = "vllm-tool-parser-py" version = "0.1.0" dependencies = [ - "criterion", - "easy-ext", - "expect-test", - "futures", - "openai-protocol", - "serde", + "pyo3", + "pythonize", "serde_json", - "thiserror 2.0.18", "thiserror-ext", - "tool-parser", - "winnow", + "vllm-parser", ] [[package]] @@ -6073,33 +5563,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "webpki-root-certs" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804f18a4ac2676ffb4e8b5b5fa9ae38af06df08162314f96a68d2a363e21a8ca" -dependencies = [ - "rustls-pki-types", -] - -[[package]] -name = "webpki-roots" -version = "0.26.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" -dependencies = [ - "webpki-roots 1.0.6", -] - -[[package]] -name = "webpki-roots" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "weezl" version = "0.1.12" @@ -6224,25 +5687,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -6260,31 +5705,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -6293,96 +5721,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" version = "1.0.2" @@ -6493,10 +5873,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" [[package]] -name = "y4m" -version = "0.8.0" +name = "xgrammar-structural-tag" +version = "0.1.0+xgrammar.0.2.2.4d145cc" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" +checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6" +dependencies = [ + "auto_impl", + "serde", + "serde_json", + "strum", + "thiserror 2.0.18", +] [[package]] name = "yoke" @@ -6632,21 +6019,40 @@ version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "zune-core" version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9" -[[package]] -name = "zune-inflate" -version = "0.2.54" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73ab332fe2f6680068f3582b16a24f90ad7096d5d39b974d1c0aff0125116f02" -dependencies = [ - "simd-adler32", -] - [[package]] name = "zune-jpeg" version = "0.5.15" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ba11bd70a53..db9764672f9 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -7,11 +7,11 @@ members = [ "src/managed-engine", "src/metrics", "src/mock-engine", - "src/reasoning-parser", + "src/parser", + "src/parser/python", "src/server", "src/text", "src/tokenizer", - "src/tool-parser", ] resolver = "3" @@ -23,46 +23,56 @@ license = "Apache-2.0" [workspace.dependencies] anyhow = "1.0.100" arc-swap = "1.9.0" -async-openai = "0.33.1" +async-openai = { version = "0.33.1", default-features = false, features = ["native-tls"] } async-trait = "0.1.89" asynk-strim-attr = "0.1.0" +auto_enums = { version = "0.8.9", features = ["tokio1"] } axum = "0.8.8" base64 = "0.22.1" bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] } byteorder = "1.5.0" -bytes = "1.11.1" +bytes = "1.12.0" clap = { version = "4.5.38", features = ["derive", "env"] } criterion = "0.5.1" easy-ext = "1.0.3" educe = "0.6.0" enum-as-inner = "0.7.0" expect-test = "1.5.1" -fastokens = "0.2.0" +fastokens = { version = "0.2.1", default-features = false } futures = "0.3.31" half = { version = "2.7.1", features = ["bytemuck"] } hex = "0.4.3" -hf-hub = { version = "0.5.0", features = ["tokio"] } +hf-hub = { version = "0.5.0", default-features = false, features = ["tokio"] } http-body = "1.0.1" +hyper = { version = "1.10.1", features = ["http1", "server"] } +hyper-util = { version = "0.1.20", features = [ + "server-graceful", + "service", + "tokio", +] } indexmap = "2.13.0" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "c8a29dcc755139fdc26185f400ea48c6d6d48273" } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } -ndarray = { version = "0.16.1", features = ["serde"] } -openai-harmony = "0.0.8" +ndarray = { version = "0.17", features = ["serde"] } +openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } openai-protocol = "1.6.0" +openssl = "0.10" parking_lot = "0.12.5" paste = "1.0.15" prometheus-client = "0.24.0" prometheus-client-derive-encode = "0.5.0" prost = "0.14.3" prost-types = "0.14.3" +pyo3 = "0.28.3" +pythonize = "0.28.0" rand = "0.9.2" reasoning-parser = "1.2.2" -reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] } +reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] } riptoken = { version = "0.3.0", default-features = false } rmp-serde = "1.3.1" rmpv = { version = "1.3.1", features = ["with-serde"] } @@ -77,6 +87,7 @@ serde_with = "3.18.0" serial_test = { version = "3.2.0", features = ["file_locks"] } sha2 = "0.10.9" socket2 = "0.6.3" +strum = { version = "0.27.2", features = ["derive"] } subenum = "1.1.3" subtle = "2.6" task-local = "0.1.1" @@ -86,6 +97,7 @@ thiserror = "2.0.16" thiserror-ext = "0.3.0" tiktoken-rs = "0.9.1" time = { version = "0.3.47", features = ["formatting", "local-offset", "macros"] } +tls-listener = { version = "0.11.2", default-features = false, features = ["openssl", "tokio-net", "axum"] } tokenizers = "0.22.0" tokio = { version = "1.47.1", features = [ "macros", @@ -94,6 +106,7 @@ tokio = { version = "1.47.1", features = [ "sync", "time", ] } +tokio-openssl = "0.6" tokio-stream = "0.1" tokio-util = { version = "0.7.18", features = ["rt"] } tonic = "0.14.5" @@ -102,7 +115,7 @@ tonic-prost = "0.14.5" tonic-prost-build = "0.14.5" tool-parser = "1.2.0" tower = { version = "0.5.3", features = ["util"] } -tower-http = { version = "0.6.8", features = ["trace"] } +tower-http = { version = "0.6.8", features = ["cors", "trace"] } tracing = { version = "0.1.44", features = ["release_max_level_debug"] } tracing-futures = { version = "0.2.5", features = ["futures-03"] } tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] } @@ -114,12 +127,12 @@ vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } -vllm-reasoning-parser = { path = "src/reasoning-parser" } +vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } -vllm-tool-parser = { path = "src/tool-parser" } -winnow = "1.0.2" +winnow = { version = "1.0.2", features = ["simd"] } +xgrammar-structural-tag = "0.1.0" zeromq = { version = "0.6.0", default-features = false, features = [ "tokio-runtime", "all-transport", @@ -131,6 +144,13 @@ too_many_arguments = "allow" [profile.dev] panic = "abort" +# Speed up cold tokenizer construction in tests. +[profile.dev.package] +fastokens = { opt-level = 3 } +regex-automata = { opt-level = 3 } +serde_json = { opt-level = 3 } +tokenizers = { opt-level = 3 } + [profile.release] lto = "thin" panic = "abort" diff --git a/rust/README.md b/rust/README.md index 679a7f0966e..b14aba3fae1 100644 --- a/rust/README.md +++ b/rust/README.md @@ -71,7 +71,7 @@ To build the `vllm-rs` in isolation: ```bash # from the local checkout -cargo install --path src/cmd --bin vllm-rs +./build_rust.sh ``` ### Example Request diff --git a/rust/deny.toml b/rust/deny.toml new file mode 100644 index 00000000000..25bd8e3831a --- /dev/null +++ b/rust/deny.toml @@ -0,0 +1,15 @@ +[bans] +multiple-versions = "allow" + +deny = [ + # TLS / crypto provider + # We prefer the system's TLS (e.g. OpenSSL) over Rust implementations. + { name = "rustls" }, + { name = "ring" }, + { name = "aws-lc-rs" }, + { name = "aws-lc-sys" }, + { name = "s2n-tls" }, + { name = "s2n-tls-sys" }, + { name = "boring" }, + { name = "boring-sys" }, +] diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index 56c5f36442d..2509d5071b6 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + syntax = "proto3"; package vllm; @@ -107,6 +110,9 @@ message KVCacheParameters { // KV Connector transfer parameters google.protobuf.Struct kv_transfer_params = 3; + + // Encoder cache connector transfer parameters + google.protobuf.Struct ec_transfer_params = 4; } // Controls which extra candidate tokens at each position should be returned @@ -173,6 +179,7 @@ message FinishInfo { google.protobuf.Struct kv_transfer_params = 6; //uint64 seed = 7; + google.protobuf.Struct ec_transfer_params = 8; } // Info for candidate tokens other than the input/sampled diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 0523b9defe9..00a6e223b77 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -21,25 +21,28 @@ serde.workspace = true serde-json-fmt.workspace = true serde_json.workspace = true serde_with.workspace = true +strum.workspace = true subenum.workspace = true thiserror.workspace = true thiserror-ext.workspace = true +time.workspace = true tokio.workspace = true tracing.workspace = true trait-set.workspace = true uuid.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true -vllm-reasoning-parser.workspace = true +vllm-parser.workspace = true vllm-text.workspace = true vllm-tokenizer.workspace = true -vllm-tool-parser.workspace = true +xgrammar-structural-tag.workspace = true [dev-dependencies] anyhow.workspace = true bytes.workspace = true clap.workspace = true expect-test.workspace = true +ndarray.workspace = true paste.workspace = true rmp-serde.workspace = true serial_test.workspace = true @@ -48,6 +51,7 @@ tokio.workspace = true tracing-subscriber.workspace = true uuid.workspace = true vllm-engine-core-client = { workspace = true, features = ["test-util"] } +vllm-tokenizer = { workspace = true, features = ["test-utils"] } zeromq.workspace = true [lints] diff --git a/rust/src/chat/examples/external_engine_chat_qwen.rs b/rust/src/chat/examples/external_engine_chat_qwen.rs index d99d672d5eb..4f629820369 100644 --- a/rust/src/chat/examples/external_engine_chat_qwen.rs +++ b/rust/src/chat/examples/external_engine_chat_qwen.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -131,13 +134,13 @@ async fn main() -> Result<()> { ChatEvent::LogprobsDelta { .. } => {} ChatEvent::Done { message, - output_token_count, + usage, finish_reason: reason, .. } => { final_reasoning = message.reasoning().unwrap_or_default(); final_text = message.text(); - final_output_token_count = output_token_count; + final_output_token_count = usage.output_token_count; finish_reason = Some(reason); break; } diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 77ed24de854..5f3a656ae04 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::sync::Arc; use tracing::info; @@ -10,12 +13,14 @@ use crate::backend::{ NewChatOutputProcessorOptions, }; use crate::error::Result; -use crate::multimodal::MultimodalModelInfo; +use crate::multimodal::{MultimodalConfigFiles, MultimodalModelInfo}; use crate::output::{ DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides, }; use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; -use crate::renderer::{DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer}; +use crate::renderer::{ + DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, +}; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -44,8 +49,12 @@ impl HfChatBackend { MultimodalModelInfo::from_paths( model_id.clone(), (!model_type.is_empty()).then_some(model_type.to_string()), - files.config_path.as_deref(), - files.preprocessor_config_path.as_deref(), + MultimodalConfigFiles { + config: files.config_path.as_deref(), + preprocessor_config: files.preprocessor_config_path.as_deref(), + video_preprocessor_config: files.video_preprocessor_config_path.as_deref(), + processor_config: files.processor_config_path.as_deref(), + }, tokenizer.clone(), )? }; @@ -61,6 +70,7 @@ impl HfChatBackend { )?), RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()), RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), + RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), }; info!( @@ -136,8 +146,11 @@ pub(super) async fn load_model_backends( fn resolve_multimodal_render_info( info: Option<&MultimodalModelInfo>, ) -> Option { + use llm_multimodal::Modality; + info.map(|info| MultimodalRenderInfo { - placeholder_token: info.placeholder_token().to_string(), + image_token: info.placeholder_token(Modality::Image).map(str::to_string), + video_token: info.placeholder_token(Modality::Video).map(str::to_string), }) } @@ -148,13 +161,16 @@ mod tests { use std::sync::Arc; use tempfile::tempdir; + use thiserror_ext::AsReport as _; + use vllm_text::Prompt; use vllm_text::backend::hf::TokenizerSource; - use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; + use vllm_text::tokenizer::DynTokenizer; + use vllm_tokenizer::test_utils::TestTokenizer; use super::HfChatBackend; - use crate::RendererSelection; - use crate::backend::{ChatBackend, LoadModelBackendsOptions}; + use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions}; use crate::request::{ChatContent, ChatMessage, ChatRequest}; + use crate::{ParserSelection, RendererSelection}; fn request_with_user_text(text: &str) -> ChatRequest { ChatRequest { @@ -186,45 +202,23 @@ mod tests { tokenizer_config_path: Some(tokenizer_config_path), generation_config_path: None, preprocessor_config_path: None, + video_preprocessor_config_path: None, + processor_config_path: None, chat_template_path: None, config_path: Some(config_path), } } - struct TestTokenizer; - - impl Tokenizer for TestTokenizer { - fn encode( - &self, - _text: &str, - _add_special_tokens: bool, - ) -> vllm_text::tokenizer::Result> { - Ok(Vec::new()) - } - - fn decode( - &self, - _token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_text::tokenizer::Result { - Ok(String::new()) - } - - fn token_to_id(&self, _token: &str) -> Option { - None - } - } - fn test_tokenizer() -> DynTokenizer { - Arc::new(TestTokenizer) + Arc::new(TestTokenizer::new()) } - fn render_prompt( + fn backend_for_selection( renderer: RendererSelection, config_json: &str, tokenizer_config_json: &str, - ) -> String { - let backend = HfChatBackend::from_resolved_model_files( + ) -> HfChatBackend { + HfChatBackend::from_resolved_model_files( resolved_files(config_json, tokenizer_config_json), "test-model".to_string(), LoadModelBackendsOptions { @@ -236,9 +230,15 @@ mod tests { }, test_tokenizer(), ) - .unwrap(); + .unwrap() + } - backend + fn render_prompt( + renderer: RendererSelection, + config_json: &str, + tokenizer_config_json: &str, + ) -> String { + backend_for_selection(renderer, config_json, tokenizer_config_json) .chat_renderer() .render(&request_with_user_text("hello")) .unwrap() @@ -272,6 +272,35 @@ mod tests { assert_eq!(prompt, "hello"); } + #[test] + fn auto_uses_harmony_renderer_and_output_processor_for_gpt_oss_model_type() { + let backend = backend_for_selection( + RendererSelection::Auto, + r#"{"model_type":"gpt_oss"}"#, + r#"{"chat_template":"{{ messages[0].content }}"}"#, + ); + + let prompt = + backend.chat_renderer().render(&request_with_user_text("hello")).unwrap().prompt; + assert!(matches!(prompt, Prompt::TokenIds(_))); + + let mut request = request_with_user_text("hello"); + let error = match backend.new_chat_output_processor( + &mut request, + NewChatOutputProcessorOptions { + tool_call_parser: &ParserSelection::Explicit("json".to_string()), + reasoning_parser: &ParserSelection::Auto, + }, + ) { + Ok(_) => panic!("gpt_oss should reject generic parser overrides"), + Err(error) => error, + }; + assert_eq!( + error.to_report_string(), + "gpt_oss uses native Harmony output parsing; generic tool parser override `json` is not supported" + ); + } + #[test] fn language_model_only_skips_multimodal_preprocessor_config() { let mut files = resolved_files( @@ -286,7 +315,7 @@ mod tests { .unwrap() .join("preprocessor_config.json"); write_json(&preprocessor_config_path, r#"{"size":[672,672]}"#); - files.preprocessor_config_path = Some(preprocessor_config_path); + files.preprocessor_config_path = Some(preprocessor_config_path.clone()); let backend = HfChatBackend::from_resolved_model_files( files.clone(), @@ -304,6 +333,9 @@ mod tests { assert!(backend.multimodal_model_info().is_none()); + let invalid_preprocessor_config = r#"{"size":[672,672]"#; + write_json(&preprocessor_config_path, invalid_preprocessor_config); + let error = HfChatBackend::from_resolved_model_files( files, "test-model".to_string(), diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index be609ba5d9e..68f2f6a568c 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::collections::HashMap; use std::sync::Arc; diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index bbd99572004..e09c3734d98 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -1,5 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use thiserror::Error; -use thiserror_ext::Macro; +use thiserror_ext::{AsReport as _, Macro}; type BoxedError = Box; @@ -18,6 +21,8 @@ pub enum Error { UnsupportedMultimodalRenderer, #[error("unsupported multimodal content: {0}")] UnsupportedMultimodalContent(&'static str), + #[error("`{modality}` input is not supported by this model")] + UnsupportedModality { modality: String }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -64,6 +69,8 @@ pub enum Error { StreamClosedBeforeTerminalOutput { request_id: String }, #[error("tool call stream state is inconsistent: {message}")] ToolCallStreamInvariant { message: String }, + #[error("failed to build structural tag: {message}")] + StructuralTag { message: String }, #[error(transparent)] Text(#[from] vllm_text::Error), #[error(transparent)] @@ -72,6 +79,45 @@ pub enum Error { pub type Result = std::result::Result; +impl Error { + /// Whether this error represents invalid user request parameters. + pub fn is_request_validation_error(&self) -> bool { + match self { + Self::PromptTooLong { .. } => true, + Self::Text(error) => error.is_request_validation_error(), + Self::UnsupportedMultimodalRenderer + | Self::UnsupportedMultimodalContent(_) + | Self::UnsupportedModality { .. } => true, + + _ => false, + } + } +} + +impl From for Error { + fn from(error: llm_multimodal::MediaConnectorError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::MultiModalError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::TransformError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::registry::ModelRegistryError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + /// Format the available-parser suffix used in user-facing error messages. fn available_parser_hint(available_names: &[String]) -> String { if available_names.is_empty() { diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index 9eb8d35042b..589fde75400 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -1,7 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::ops::Deref; use std::sync::Arc; use serde::{Deserialize, Serialize}; +use vllm_llm::TokenUsage; use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs}; use crate::FinishReason; @@ -197,13 +201,12 @@ pub enum ChatEvent { /// metadata. Done { message: AssistantMessage, - /// Number of prompt tokens actually sent to the engine after chat - /// template rendering and tokenization. - prompt_token_count: usize, - /// Number of output tokens generated. - output_token_count: usize, + usage: TokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for + /// disaggregated serving. + ec_transfer_params: Option, }, } diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index ae366a4de05..625a2d7818b 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Minimal chat facade above [`vllm_text`]. //! //! This crate keeps the northbound boundary intentionally small: @@ -29,8 +32,8 @@ pub use parser::reasoning::{ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ - ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, RenderedPrompt, - RendererSelection, + ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, + HarmonyChatRenderer, RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -50,7 +53,8 @@ mod request; mod stream; use vllm_engine_core_client::EngineCoreClient; -use vllm_engine_core_client::protocol::ModelDtype; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::Llm; use vllm_text::{Prompt, TextLlm, TextRequest}; @@ -140,6 +144,16 @@ impl ChatLlm { self } + /// Tokenizer vocabulary size. + pub fn tokenizer_vocab_size(&self) -> usize { + self.text.tokenizer_vocab_size() + } + + /// Model vocabulary size from the model config. + pub fn model_vocab_size(&self) -> usize { + self.text.model_vocab_size() + } + /// Expose the underlying text facade for raw text-generation routes such as /// `/v1/completions`. pub fn text(&self) -> &TextLlm { @@ -161,6 +175,9 @@ impl ChatLlm { pub async fn chat(&self, mut request: ChatRequest) -> Result { request.validate()?; + // Stamp before rendering so render and tokenize count toward TTFT/e2e. + let arrival_time = vllm_llm::current_unix_timestamp_secs(); + let output_processor = self.backend.new_chat_output_processor( &mut request, NewChatOutputProcessorOptions { @@ -169,6 +186,14 @@ impl ChatLlm { }, )?; let rendered = self.backend.chat_renderer().render(&request)?; + let reasoning_parser_kwargs = + request + .sampling_params + .structured_outputs + .is_some() + .then(|| ReasoningParserKwargs { + chat_template_kwargs: rendered.effective_template_kwargs.clone(), + }); let (prompt, mm_features) = multimodal::finalize_rendered_prompt( &request, @@ -189,7 +214,9 @@ impl ChatLlm { cache_salt: request.cache_salt, add_special_tokens: request.add_special_tokens, data_parallel_rank: request.data_parallel_rank, + reasoning_parser_kwargs, lora_request: request.lora_request, + arrival_time: Some(arrival_time), }; let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); @@ -225,6 +252,12 @@ impl ChatLlm { Ok(token_ids) } + /// Abort in-flight requests by their external (user-supplied) request ids. + pub async fn abort(&self, external_ids: &[String]) -> Result<()> { + self.text.abort(external_ids).await?; + Ok(()) + } + /// Shut down the underlying LLM client and its background tasks. pub async fn shutdown(self) -> Result<()> { self.text.shutdown().await?; @@ -261,7 +294,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string()); } #[test] @@ -272,6 +305,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, step3)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index fcfee0ccb33..662e6fd0649 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -1,8 +1,11 @@ -//! Chat-layer multimodal image preparation. +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Chat-layer multimodal media preparation. //! -//! This module owns the narrow image-only multimodal path for chat requests: -//! it extracts image parts from structured chat messages, fetches and -//! preprocesses them through `llm-multimodal`, expands rendered prompt +//! This module owns the multimodal path for chat requests: it extracts media +//! parts from structured chat messages, fetches and preprocesses them through +//! `llm-multimodal` one modality at a time, expands rendered prompt //! placeholders after tokenization, and builds the engine-facing //! `MmFeatures` payload. //! @@ -16,18 +19,15 @@ use std::sync::{Arc, LazyLock}; use itertools::izip; use llm_multimodal::{ - AsyncMultiModalTracker, FieldLayout, ImagePreProcessor, ImageProcessorRegistry, MediaConnector, - MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec, - ModelRegistry, PreProcessorConfig, PreprocessedImages, PromptReplacement, TokenResolver, - TrackedMedia, + AsyncMultiModalTracker, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, + MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, + PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, Tokenizer as TokenResolver, + TrackedMedia, VideoClip, VisionPreProcessor, VisionProcessorRegistry, }; +use thiserror_ext::AsReport as _; use tracing::warn; -use vllm_engine_core_client::protocol::ModelDtype; -use vllm_engine_core_client::protocol::multimodal::{ - MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem, - MmSharedField, MmSlice, PlaceholderRange, SliceSpec, -}; -use vllm_engine_core_client::protocol::tensor::WireTensor; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem}; use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; @@ -35,14 +35,20 @@ use crate::error::{Error, Result, bail_multimodal, multimodal}; use crate::renderer::RenderedPrompt; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; +mod expand; +mod image; mod tensor; +mod video; + +use self::expand::expand_prompt_token_ids; /// Resolved multimodal support for one loaded model. #[derive(Clone)] pub struct MultimodalModelInfo { context: MultimodalModelContext, spec: ResolvedMultimodalSpec, - image_processor: ResolvedImageProcessor, + image: Option, + video: Option, media_connector: Arc, } @@ -74,91 +80,171 @@ impl MultimodalModelContext { REGISTRY.lookup(&self.metadata()) } - /// Resolve a static image preprocessor for one loaded model. - fn resolve_image_processor(&self) -> Option<&'static dyn ImagePreProcessor> { - static REGISTRY: LazyLock = - LazyLock::new(ImageProcessorRegistry::with_defaults); + /// Resolve a static vision preprocessor for one loaded model. + /// + /// The vision preprocessor serves both the image and video modalities. + fn resolve_vision_processor(&self) -> Option<&'static dyn VisionPreProcessor> { + static REGISTRY: LazyLock = + LazyLock::new(VisionProcessorRegistry::with_defaults); REGISTRY.find(&self.model_id, self.model_type.as_deref()) } } -/// Static model-specific prompt and tensor-layout behavior. +/// Static model-specific tensor-layout behavior shared across modalities. #[derive(Clone)] struct ResolvedMultimodalSpec { raw: &'static dyn ModelProcessorSpec, - placeholder_token: String, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, field_layouts: HashMap, keep_on_cpu_keys: HashSet, } impl ResolvedMultimodalSpec { - fn new(raw: &'static dyn ModelProcessorSpec, context: &MultimodalModelContext) -> Result { - let metadata = context.metadata(); - let placeholder_token = - raw.placeholder_token(&metadata).map_err(|error| multimodal!("{error}"))?; - // This is the rendered prompt marker, so resolve it from the token - // string itself. Do not use `ModelProcessorSpec::placeholder_token_id()`: - // for some specs that ID is the replacement vision/patch token, - // not necessarily the token ID of `placeholder_token`. - let placeholder_marker_token_id = - context.tokenizer().token_to_id(&placeholder_token).ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` is not in the tokenizer vocabulary" - ) - })?; - let placeholder_embed_token_id = - raw.placeholder_token_id(&metadata).map_err(|error| multimodal!("{error}"))? as u32; - - Ok(Self { + fn new(raw: &'static dyn ModelProcessorSpec) -> Self { + Self { raw, - placeholder_token, - placeholder_marker_token_id, - placeholder_embed_token_id, field_layouts: raw.field_layouts(), keep_on_cpu_keys: raw.keep_on_cpu_keys().into_iter().collect(), - }) + } } - fn prompt_replacements( + fn prompt_replacements_for( &self, context: &MultimodalModelContext, - preprocessed: &PreprocessedImages, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, ) -> Result> { - self.raw - .prompt_replacements(&context.metadata(), preprocessed) - .map_err(|error| multimodal!("{error}")) + Ok(self.raw.prompt_replacements_for(&context.metadata(), preprocessed, modality)?) } } -/// Static image preprocessor plus its loaded config. +/// Resolved placeholder tokens for one modality. #[derive(Clone)] -struct ResolvedImageProcessor { - raw: &'static dyn ImagePreProcessor, +struct ResolvedPlaceholder { + token: String, + /// The token ID emitted for `token` in the rendered prompt. + marker_token_id: u32, + /// The model-declared embed token ID marked in `is_embed` masks. + embed_token_id: u32, +} + +impl ResolvedPlaceholder { + fn resolve( + raw: &'static dyn ModelProcessorSpec, + context: &MultimodalModelContext, + modality: Modality, + ) -> Result { + let metadata = context.metadata(); + let token = raw.placeholder_token_for(&metadata, modality)?; + // This is the rendered prompt marker, so resolve it from the token + // string itself. Do not use `ModelProcessorSpec::placeholder_token_id_for()`: + // for some specs that ID is the replacement vision/patch token, + // not necessarily the token ID of the placeholder token. + let marker_token_id = context.tokenizer().token_to_id(&token).ok_or_else(|| { + multimodal!("placeholder token `{token}` is not in the tokenizer vocabulary") + })?; + let embed_token_id = raw.placeholder_token_id_for(&metadata, modality)? as u32; + + Ok(Self { + token, + marker_token_id, + embed_token_id, + }) + } +} + +/// Static per-modality vision preprocessor plus its loaded config and +/// resolved placeholder tokens. +#[derive(Clone)] +struct ModalitySupport { + placeholder: ResolvedPlaceholder, + processor: &'static dyn VisionPreProcessor, config: PreProcessorConfig, } -/// Request-scoped fetched media, kept together with tracker UUID metadata. -struct FetchedImageMedia { - frames: Vec>, - uuids: Vec>, +/// Model-repo config file locations consumed by multimodal support. +#[derive(Debug, Default, Clone, Copy)] +pub struct MultimodalConfigFiles<'a> { + pub config: Option<&'a Path>, + pub preprocessor_config: Option<&'a Path>, + /// Video-specific preprocessor config (`video_preprocessor_config.json`). + pub video_preprocessor_config: Option<&'a Path>, + /// Combined processor config (`processor_config.json`), whose modality + /// sections are fallback preprocessor config sources. + pub processor_config: Option<&'a Path>, +} + +/// Load a modality's dedicated preprocessor config, falling back to its section +/// in the combined processor config. +fn load_preprocessor_config( + dedicated_path: Option<&Path>, + dedicated_name: &str, + processor_config_path: Option<&Path>, + processor_section: &str, +) -> Result> { + if let Some(path) = dedicated_path { + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read {dedicated_name}: {error}"))?; + let config = PreProcessorConfig::from_json(&text) + .map_err(|error| multimodal!("failed to parse {dedicated_name}: {error}"))?; + return Ok(Some(config)); + } + + let Some(path) = processor_config_path else { + return Ok(None); + }; + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read processor_config.json: {error}"))?; + let value: serde_json::Value = serde_json::from_str(&text) + .map_err(|error| multimodal!("failed to parse processor_config.json: {error}"))?; + let Some(processor) = value.get(processor_section) else { + return Ok(None); + }; + let config = PreProcessorConfig::from_value(processor.clone()).map_err(|error| { + multimodal!("failed to parse {processor_section} from processor_config.json: {error}") + })?; + Ok(Some(config)) +} + +/// Request-scoped fetched media, split per modality with tracker UUID +/// metadata preserved in request order. +struct FetchedMedia { + images: Vec>, + image_uuids: Vec>, + videos: Vec>, + video_uuids: Vec>, +} + +/// One modality's preprocessed output, ready for the shared expansion and +/// feature-assembly tail. +struct PreparedMedia { + modality: Modality, + placeholder: ResolvedPlaceholder, + /// One replacement per media item, in request order. + replacements: Vec, + /// One entry per media item, aligned with `replacements`. + items: Vec, +} + +/// One media item's complete engine kwargs plus identity metadata. +struct PreparedItem { + data: MmKwargsItem, + hash: String, + uuid: Option, } impl MultimodalModelInfo { /// Load and resolve multimodal support from model files. /// - /// Returns `Ok(Some(_))` only when both the model spec and image processor - /// are registered. File read/parse failures are real errors; unsupported - /// model families are logged and returned as `Ok(None)`. + /// Returns `Ok(Some(_))` only when the model spec is registered and at + /// least one modality resolves. File read/parse failures are real errors; + /// unsupported model families are logged and returned as `Ok(None)`. pub fn from_paths( model_id: String, model_type: Option, - config_path: Option<&Path>, - preprocessor_config_path: Option<&Path>, + files: MultimodalConfigFiles<'_>, tokenizer: DynTokenizer, ) -> Result> { - let config = match config_path { + let config = match files.config { Some(path) => { let text = fs::read_to_string(path) .map_err(|error| multimodal!("failed to read config.json: {error}"))?; @@ -167,17 +253,20 @@ impl MultimodalModelInfo { } None => serde_json::Value::Object(Default::default()), }; - let preprocessor_config = match preprocessor_config_path { - Some(path) => { - let text = fs::read_to_string(path).map_err(|error| { - multimodal!("failed to read preprocessor_config.json: {error}") - })?; - PreProcessorConfig::from_json(&text).map_err(|error| { - multimodal!("failed to parse preprocessor_config.json: {error}") - })? - } - None => PreProcessorConfig::default(), - }; + let image_preprocessor_config = load_preprocessor_config( + files.preprocessor_config, + "preprocessor_config.json", + files.processor_config, + "image_processor", + )? + .unwrap_or_default(); + let video_preprocessor_config = load_preprocessor_config( + files.video_preprocessor_config, + "video_preprocessor_config.json", + files.processor_config, + "video_processor", + )? + .unwrap_or_else(|| image_preprocessor_config.clone()); let context = MultimodalModelContext { model_id, @@ -186,7 +275,21 @@ impl MultimodalModelInfo { tokenizer: TokenizerResolver(tokenizer), }; - let Some(spec) = context.resolve_model_spec() else { + Self::from_loaded( + context, + image_preprocessor_config, + video_preprocessor_config, + ) + } + + /// Resolve multimodal support from an assembled context and parsed + /// preprocessor configs. + fn from_loaded( + context: MultimodalModelContext, + image_preprocessor_config: PreProcessorConfig, + video_preprocessor_config: PreProcessorConfig, + ) -> Result> { + let Some(raw_spec) = context.resolve_model_spec() else { warn!( model_id = context.model_id, model_type = context.model_type, @@ -194,47 +297,99 @@ impl MultimodalModelInfo { ); return Ok(None); }; - let spec = ResolvedMultimodalSpec::new(spec, &context)?; - let Some(image_processor) = context.resolve_image_processor() else { + let Some(processor) = context.resolve_vision_processor() else { warn!( model_id = context.model_id, model_type = context.model_type, - "image processor is not registered; disabling multimodal support for this model" + "vision processor is not registered; disabling multimodal support for this model" ); return Ok(None); }; - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()) - .map_err(|error| multimodal!("{error}"))?, - ); + // Warn and disable the modality if the placeholder resolution fails. + let resolve_placeholder = + |modality: Modality| match ResolvedPlaceholder::resolve(raw_spec, &context, modality) { + Ok(placeholder) => Some(placeholder), + Err(error) => { + warn!( + model_id = context.model_id, + %modality, + error = %error.as_report(), + "placeholder tokens did not resolve; disabling this modality for this model" + ); + None + } + }; + + let image = resolve_placeholder(Modality::Image).map(|placeholder| ModalitySupport { + placeholder, + processor, + config: image_preprocessor_config, + }); + + let video = resolve_placeholder(Modality::Video).and_then(|placeholder| { + // Placeholder expansion attributes markers to modalities by token + // ID, so a marker shared with the image modality is ambiguous. + let image_marker = image.as_ref().map(|image| image.placeholder.marker_token_id); + if image_marker == Some(placeholder.marker_token_id) { + warn!( + model_id = context.model_id, + token = placeholder.token, + "video placeholder token collides with the image placeholder; disabling video support for this model" + ); + None + } else { + Some(ModalitySupport { + placeholder, + processor, + config: video_preprocessor_config, + }) + } + }); + + if image.is_none() && video.is_none() { + warn!( + model_id = context.model_id, + model_type = context.model_type, + "no multimodal modality resolved; disabling multimodal support for this model" + ); + return Ok(None); + } + + let media_connector = Arc::new(MediaConnector::new( + reqwest::Client::new(), + MediaConnectorConfig::default(), + )?); Ok(Some(Self { context, - spec, - image_processor: ResolvedImageProcessor { - raw: image_processor, - config: preprocessor_config, - }, + spec: ResolvedMultimodalSpec::new(raw_spec), + image, + video, media_connector, })) } - /// Return the template-visible placeholder token for this model. + /// Return the template-visible placeholder token for one modality, when + /// this model supports it. /// - /// The HF renderer uses this token while flattening image content in string - /// content format. - pub(crate) fn placeholder_token(&self) -> &str { - &self.spec.placeholder_token + /// The HF renderer uses these tokens while flattening media content in + /// string content format. + pub fn placeholder_token(&self, modality: Modality) -> Option<&str> { + match modality { + Modality::Image => self.image.as_ref()?.placeholder.token.as_str().into(), + Modality::Video => self.video.as_ref()?.placeholder.token.as_str().into(), + _ => None, + } } } /// Finalize a rendered chat prompt into text-generation input. /// /// Text-only requests pass through unchanged as `Prompt::Text`. Multimodal -/// requests are tokenized in chat, their image placeholders are expanded, and -/// preprocessed image features are attached for engine-core transport. +/// requests are tokenized in chat, their media placeholders are expanded, and +/// preprocessed media features are attached for engine-core transport. pub(crate) async fn finalize_rendered_prompt( request: &ChatRequest, rendered: RenderedPrompt, @@ -245,22 +400,21 @@ pub(crate) async fn finalize_rendered_prompt( return Ok((rendered.prompt, None)); } let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?; - let Prompt::Text(prompt) = rendered.prompt else { - bail_multimodal!("multimodal chat renderer must return a text prompt before expansion"); + let mut prompt_token_ids = match rendered.prompt { + Prompt::Text(prompt) => info + .context + .tokenizer() + .encode(&prompt, request.add_special_tokens) + .map_err(|error| multimodal!("{error}"))?, + Prompt::TokenIds(token_ids) => token_ids, }; let media_parts = extract_media_parts(request)?; - - let mut prompt_token_ids = info - .context - .tokenizer() - .encode(&prompt, request.add_special_tokens) - .map_err(|error| multimodal!("{error}"))?; let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?; Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract image media parts from chat messages in message/content order. +/// Extract media parts from chat messages in message/content order. /// /// Assistant history is skipped because generated assistant blocks are already /// represented as text for prompt rendering in this crate. @@ -289,6 +443,12 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { detail: *detail, uuid: uuid.clone(), }), + ChatContentPart::VideoUrl { video_url, uuid } => { + all_parts.push(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }) + } } } } @@ -296,8 +456,8 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { } impl MultimodalModelInfo { - /// Run media fetch, image preprocessing, prompt expansion, and feature - /// build. + /// Run media fetch, per-modality preprocessing, prompt expansion, and + /// feature build. /// /// `prompt_token_ids` is mutated in place because placeholder expansion /// changes both the final prompt and the offsets recorded in @@ -313,12 +473,47 @@ impl MultimodalModelInfo { } let media_parts_len = media_parts.len(); - let fetched = self.fetch_images(media_parts).await?; - let preprocessed = self.preprocess_images(&fetched.frames).await?; - let replacements = self.spec.prompt_replacements(&self.context, &preprocessed)?; - let ranges = self.expand_prompt_tokens(prompt_token_ids, replacements)?; + // TODO: enforce per-modality item-count limits, aligned with the + // engine's `--limit-mm-per-prompt` semantics. + let fetched = self.fetch_media(media_parts).await?; + + let mut prepared = Vec::new(); + if !fetched.images.is_empty() { + prepared + .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?); + } + if !fetched.videos.is_empty() { + prepared + .push(self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype).await?); + } + + let mut ranges = expand_prompt_token_ids(prompt_token_ids, &prepared)?; + + let mut features = Vec::with_capacity(media_parts_len); + for media in prepared { + let media_ranges = ranges.remove(&media.modality).unwrap_or_default(); + if media_ranges.len() != media.items.len() { + bail_multimodal!( + "number of expanded `{}` placeholders {} does not match number of media items {}", + media.modality, + media_ranges.len(), + media.items.len() + ); + } + for (item, range) in izip!(media.items, media_ranges) { + features.push(MmFeatureSpec { + data: Some(item.data), + modality: media.modality.to_string(), + identifier: item.uuid.unwrap_or_else(|| item.hash.clone()), + mm_position: range, + mm_hash: Some(item.hash), + }); + } + } + // Mirror the Python frontend (`argsort_mm_positions`): features are + // ordered by their placeholder position in the prompt. + features.sort_by_key(|feature| feature.mm_position.offset); - let features = self.build_features(preprocessed, fetched, ranges, model_dtype)?; if features.len() != media_parts_len { bail_multimodal!( "number of built multimodal features {} does not match number of media parts {}", @@ -329,218 +524,51 @@ impl MultimodalModelInfo { Ok(features) } - /// Fetch all image parts and preserve their request-order UUID metadata. - async fn fetch_images(&self, media_parts: Vec) -> Result { + /// Fetch all media parts and split them per modality, preserving their + /// request-order UUID metadata. + async fn fetch_media(&self, media_parts: Vec) -> Result { let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector)); for part in media_parts { - tracker.push_part(part).map_err(|error| multimodal!("{error}"))?; + tracker.push_part(part)?; } - let tracker_output = tracker.finalize().await.map_err(|error| multimodal!("{error}"))?; - let images = tracker_output.data.get(&Modality::Image).cloned().unwrap_or_default(); - let uuids = tracker_output.uuids.get(&Modality::Image).cloned().unwrap_or_default(); + let mut tracker_output = tracker.finalize().await?; - let frames = images + let images = tracker_output + .data + .remove(&Modality::Image) + .unwrap_or_default() .into_iter() .map(|media| match media { TrackedMedia::Image(frame) => Ok(frame), - _ => Err(Error::UnsupportedMultimodalContent("non-image")), + _ => Err(multimodal!( + "tracker returned non-image media for the image modality" + )), }) .collect::>>()?; + let image_uuids = tracker_output.uuids.remove(&Modality::Image).unwrap_or_default(); - Ok(FetchedImageMedia { frames, uuids }) - } + let videos = tracker_output + .data + .remove(&Modality::Video) + .unwrap_or_default() + .into_iter() + .map(|media| match media { + TrackedMedia::Video(clip) => Ok(clip), + _ => Err(multimodal!( + "tracker returned non-video media for the video modality" + )), + }) + .collect::>>()?; + let video_uuids = tracker_output.uuids.remove(&Modality::Video).unwrap_or_default(); - /// Preprocess fetched image frames with the model's resolved image - /// processor. - /// - /// The processor work is CPU-heavy relative to request wiring, so it runs - /// in a blocking task and returns owned tensors ready for wire - /// conversion. - async fn preprocess_images( - &self, - image_frames: &[Arc], - ) -> Result { - let config = self.image_processor.config.clone(); - let processor = self.image_processor.raw; - let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); - - tokio::task::spawn_blocking(move || { - processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}")) + Ok(FetchedMedia { + images, + image_uuids, + videos, + video_uuids, }) - .await - .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? } - - /// Replace rendered placeholder markers with model-specific replacement - /// tokens. - /// - /// Replacements are consumed in order, matching the original media-part - /// order. The returned ranges point into the already-expanded prompt. - fn expand_prompt_tokens( - &self, - prompt_token_ids: &mut Vec, - replacements: Vec, - ) -> Result> { - expand_prompt_token_ids( - prompt_token_ids, - replacements, - self.spec.placeholder_marker_token_id, - self.spec.placeholder_embed_token_id, - &self.spec.placeholder_token, - ) - } - - /// Convert preprocessed image tensors into engine-core multimodal features. - /// - /// One `MmFeatureSpec` is produced per image. Tensor fields are - /// sliced according to the model spec's field layout declarations. - fn build_features( - &self, - preprocessed: PreprocessedImages, - images: FetchedImageMedia, - ranges: Vec, - model_dtype: ModelDtype, - ) -> Result { - let len = images.frames.len(); - let tensors = tensor::collect_tensors(preprocessed, model_dtype)?; - - let mut features = Vec::with_capacity(images.frames.len()); - for (index, (frame, uuid, range)) in izip!(images.frames, images.uuids, ranges).enumerate() - { - let mut data = MmKwargsItem::new(); - for (key, tensor) in &tensors { - let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); - let (value, field) = match self.spec.field_layouts.get(key) { - Some(FieldLayout::Batched) => ( - tensor.batched_value_at(index)?, - MmField::Batched(MmBatchedField { keep_on_cpu }), - ), - Some(FieldLayout::Flat { sizes_key }) => { - let sizes = tensors.get(sizes_key).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` is missing") - })?; - let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; - ( - tensor.flat_value_range(start, end)?, - MmField::Flat(MmFlatField { - slices: vec![MmSlice::Slice(SliceSpec { - start: Some(0), - stop: Some((end - start) as isize), - step: None, - })], - dim: 0, - keep_on_cpu, - }), - ) - } - None => ( - tensor.clone(), - MmField::Shared(MmSharedField { - batch_size: len, - keep_on_cpu, - }), - ), - }; - - data.insert( - key.clone(), - MmFieldElem { - data: Some(value.try_into()?), - field, - }, - ); - } - - let hash = frame.hash.clone(); - features.push(MmFeatureSpec { - data: Some(data), - modality: "image".to_string(), - identifier: uuid.unwrap_or_else(|| hash.clone()), - mm_position: range, - mm_hash: Some(hash), - }); - } - - Ok(features) - } -} - -fn expand_prompt_token_ids( - prompt_token_ids: &mut Vec, - replacements: Vec, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, - placeholder_token: &str, -) -> Result> { - if replacements.is_empty() { - return Ok(Vec::new()); - } - - let replacement_growth = replacements.iter().fold(0usize, |total, replacement| { - total.saturating_add(replacement.tokens.len().saturating_sub(1)) - }); - let mut expanded = - Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth)); - let mut ranges = Vec::with_capacity(replacements.len()); - let mut cursor = 0usize; - - for replacement in replacements { - if replacement.modality != Modality::Image { - bail_multimodal!( - "unsupported prompt replacement modality `{}`", - replacement.modality - ); - } - - let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor) - .ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` was not found in tokenized prompt" - ) - })?; - - if replacement.tokens.is_empty() { - bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens"); - } - - let replacement_len = replacement.tokens.len(); - let is_embed = { - let mask = replacement - .tokens - .iter() - .map(|&token| token as u32 == placeholder_embed_token_id) - .collect::>(); - WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? - }; - - expanded.extend_from_slice(&prompt_token_ids[cursor..offset]); - let expanded_offset = expanded.len(); - expanded.extend(replacement.tokens.into_iter().map(|token| token as u32)); - ranges.push(PlaceholderRange { - offset: expanded_offset, - length: replacement_len, - is_embed: Some(is_embed), - }); - cursor = offset + 1; - } - - expanded.extend_from_slice(&prompt_token_ids[cursor..]); - *prompt_token_ids = expanded; - - Ok(ranges) -} - -/// Find `needle` in `haystack`, starting at `start`. -/// -/// This is intentionally order-preserving rather than a global replace: each -/// image consumes the next placeholder occurrence. -fn find_next_token(haystack: &[u32], needle: u32, start: usize) -> Option { - haystack - .get(start..)? - .iter() - .position(|token| *token == needle) - .map(|offset| start + offset) } /// Adapter from the frontend tokenizer trait to `llm-multimodal`. @@ -555,108 +583,65 @@ impl TokenResolver for TokenizerResolver { fn id_to_token(&self, id: u32) -> Option { self.0.id_to_token(id) } + + fn encode_text(&self, text: &str) -> Option> { + self.0.encode(text, false).ok() + } } #[cfg(test)] mod tests { use std::sync::Arc; - use llm_multimodal::TokenId; - use vllm_engine_core_client::protocol::tensor::WireArrayData; - use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError}; + use vllm_tokenizer::test_utils::TestTokenizer; use super::*; - const LLAMA4_IMAGE_START_ID: u32 = 200088; - const LLAMA4_IMAGE_END_ID: u32 = 200089; - const LLAMA4_IMAGE_ID: u32 = 200090; - const LLAMA4_PATCH_ID: u32 = 200092; - const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; - const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; + pub(super) const LLAMA4_IMAGE_START_ID: u32 = 200088; + pub(super) const LLAMA4_IMAGE_END_ID: u32 = 200089; + pub(super) const LLAMA4_IMAGE_ID: u32 = 200090; + pub(super) const LLAMA4_PATCH_ID: u32 = 200092; + pub(super) const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; + pub(super) const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; - struct TestTokenizer; + pub(super) const QWEN3_IMAGE_PAD_ID: u32 = 151655; + pub(super) const QWEN3_VIDEO_PAD_ID: u32 = 151656; - impl Tokenizer for TestTokenizer { - fn encode( - &self, - text: &str, - _add_special_tokens: bool, - ) -> std::result::Result, TokenizerError> { - Ok(match text { - "<|image|>" => vec![LLAMA4_IMAGE_ID], - text => text.bytes().map(u32::from).collect(), - }) - } - - fn decode( - &self, - _token_ids: &[u32], - _skip_special_tokens: bool, - ) -> std::result::Result { - Ok(String::new()) - } - - fn token_to_id(&self, token: &str) -> Option { - match token { - "<|image_start|>" => Some(LLAMA4_IMAGE_START_ID), - "<|image_end|>" => Some(LLAMA4_IMAGE_END_ID), - "<|image|>" => Some(LLAMA4_IMAGE_ID), - "<|patch|>" => Some(LLAMA4_PATCH_ID), - "<|tile_x_separator|>" => Some(LLAMA4_TILE_X_SEPARATOR_ID), - "<|tile_y_separator|>" => Some(LLAMA4_TILE_Y_SEPARATOR_ID), - _ => None, - } - } - - fn id_to_token(&self, id: u32) -> Option { - match id { - LLAMA4_IMAGE_START_ID => Some("<|image_start|>".to_string()), - LLAMA4_IMAGE_END_ID => Some("<|image_end|>".to_string()), - LLAMA4_IMAGE_ID => Some("<|image|>".to_string()), - LLAMA4_PATCH_ID => Some("<|patch|>".to_string()), - LLAMA4_TILE_X_SEPARATOR_ID => Some("<|tile_x_separator|>".to_string()), - LLAMA4_TILE_Y_SEPARATOR_ID => Some("<|tile_y_separator|>".to_string()), - _ => None, - } - } - - fn create_decode_stream( - &self, - _prompt_token_ids: &[u32], - _skip_special_tokens: bool, - _min_bytes_to_buffer: usize, - ) -> Box { - unreachable!("not used") - } + fn llama4_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|image_start|>", LLAMA4_IMAGE_START_ID) + .with_regular_token("<|image_end|>", LLAMA4_IMAGE_END_ID) + .with_regular_token("<|image|>", LLAMA4_IMAGE_ID) + .with_regular_token("<|patch|>", LLAMA4_PATCH_ID) + .with_regular_token("<|tile_x_separator|>", LLAMA4_TILE_X_SEPARATOR_ID) + .with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID) } - fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo { + pub(super) fn qwen3_vl_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|image_pad|>", QWEN3_IMAGE_PAD_ID) + .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID) + } + + fn test_info( + model_type: &str, + config: serde_json::Value, + tokenizer: TestTokenizer, + ) -> MultimodalModelInfo { let context = MultimodalModelContext { model_id: format!("{model_type}-test"), model_type: Some(model_type.to_string()), config, - tokenizer: TokenizerResolver(Arc::new(TestTokenizer)), + tokenizer: TokenizerResolver(Arc::new(tokenizer)), }; - let spec = context - .resolve_model_spec() - .unwrap_or_else(|| panic!("{model_type} spec should match")); - let spec = ResolvedMultimodalSpec::new(spec, &context).unwrap(); - let raw_image_processor = context - .resolve_image_processor() - .unwrap_or_else(|| panic!("{model_type} image processor should match")); - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()).unwrap(), - ); - MultimodalModelInfo { + MultimodalModelInfo::from_loaded( context, - spec, - image_processor: ResolvedImageProcessor { - raw: raw_image_processor, - config: PreProcessorConfig::default(), - }, - media_connector, - } + PreProcessorConfig::default(), + PreProcessorConfig::default(), + ) + .unwrap() + .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve")) } fn llama4_info() -> MultimodalModelInfo { @@ -665,173 +650,96 @@ mod tests { "image_token_index": LLAMA4_PATCH_ID, "vision_config": {"image_size": 336, "patch_size": 14} }); - test_info("llama4", config) + test_info("llama4", config, llama4_tokenizer()) } - fn llama4_single_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], + pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + "vision_config": {"patch_size": 16} + }); + test_info("qwen3_vl", config, qwen3_vl_tokenizer()) + } + + #[test] + fn from_paths_resolves_image_config_from_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + }) + .to_string(), ) - } - - fn llama4_multi_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_X_SEPARATOR_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], + .unwrap(); + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"image_processor":{"size":{"shortest_edge":64}}}"#, ) - } + .unwrap(); - fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { - let tensor = range.is_embed.as_ref().expect("is_embed mask"); - assert_eq!(tensor.dtype, "bool"); - assert_eq!(tensor.shape, vec![expected.len()]); - assert_eq!( - tensor.data, - WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) - ); + let info = MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }, + Arc::new(qwen3_vl_tokenizer()), + ) + .unwrap() + .unwrap(); + + assert_eq!(info.image.unwrap().config.get_shortest_edge(), Some(64)); } #[test] - fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let replacements = vec![llama4_multi_tile_replacement()]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); + fn qwen3_vl_resolves_image_and_video_support() { + let info = qwen3_vl_info(); assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_X_SEPARATOR_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_Y_SEPARATOR_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - ] + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 8); - assert_bool_mask( - &ranges[0], - &[false, true, false, true, false, false, true, false], + assert_eq!( + info.placeholder_token(Modality::Video), + Some("<|video_pad|>") + ); + assert_ne!( + info.image.as_ref().unwrap().placeholder.marker_token_id, + info.video.as_ref().unwrap().placeholder.marker_token_id, ); } #[test] - fn expand_prompt_tokens_errors_when_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, 2, 3]; - let replacements = vec![llama4_single_tile_replacement()]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - } - - #[test] - fn expand_prompt_tokens_ignores_empty_replacements() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap(); - - assert!(ranges.is_empty()); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } - - #[test] - fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } - - #[test] - fn expand_prompt_tokens_errors_when_replacement_is_empty() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![PromptReplacement::sequence( - Modality::Image, - "<|image|>", - Vec::new(), - )]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!( - matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) - ); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } - - #[test] - fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); + fn qwen3_vl_without_video_token_id_disables_video_support_only() { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "vision_config": {"patch_size": 16} + }); + let info = test_info("qwen3_vl", config, qwen3_vl_tokenizer()); assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 3, - ] + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 5); - assert_bool_mask(&ranges[0], &[false, false, true, true, false]); - assert_eq!(ranges[1].offset, 7); - assert_eq!(ranges[1].length, 5); - assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + assert_eq!(info.placeholder_token(Modality::Video), None); + } + + #[test] + fn llama4_resolves_image_support_only() { + let info = llama4_info(); + + assert_eq!(info.placeholder_token(Modality::Image), Some("<|image|>")); + assert_eq!(info.placeholder_token(Modality::Video), None); } } diff --git a/rust/src/chat/src/multimodal/expand.rs b/rust/src/chat/src/multimodal/expand.rs new file mode 100644 index 00000000000..b679f13b079 --- /dev/null +++ b/rust/src/chat/src/multimodal/expand.rs @@ -0,0 +1,449 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Prompt placeholder expansion shared across modalities. + +use std::collections::{HashMap, VecDeque}; + +use llm_multimodal::{Modality, PromptReplacement}; +use vllm_engine_core_client::protocol::multimodal::PlaceholderRange; +use vllm_engine_core_client::protocol::tensor::WireTensor; + +use super::PreparedMedia; +use crate::error::{Error, Result, bail_multimodal}; + +/// One modality's queue of pending placeholder replacements for prompt +/// expansion. +struct ExpansionLane<'a> { + modality: Modality, + marker_token_id: u32, + embed_token_id: u32, + placeholder_token: String, + replacements: VecDeque<&'a PromptReplacement>, +} + +impl<'a> ExpansionLane<'a> { + fn from_prepared(media: &'a PreparedMedia) -> Option { + if media.replacements.is_empty() { + return None; + } + + Some(Self { + modality: media.modality, + marker_token_id: media.placeholder.marker_token_id, + embed_token_id: media.placeholder.embed_token_id, + placeholder_token: media.placeholder.token.clone(), + replacements: media.replacements.iter().collect(), + }) + } +} + +/// Replace rendered placeholder markers with model-specific replacement +/// tokens across all modalities in one left-to-right pass. +/// +/// Each prepared modality consumes its own marker occurrences in order, +/// matching the original media-part order within that modality; markers of +/// different modalities may interleave freely. +/// +/// The returned ranges point into the already-expanded prompt, grouped per +/// modality in item order. +pub(super) fn expand_prompt_token_ids( + prompt_token_ids: &mut Vec, + prepared: &[PreparedMedia], +) -> Result>> { + let mut lanes = prepared.iter().filter_map(ExpansionLane::from_prepared).collect::>(); + if lanes.is_empty() { + return Ok(HashMap::new()); + } + + let replacement_growth = lanes + .iter() + .flat_map(|lane| lane.replacements.iter()) + .fold(0usize, |total, replacement| { + total.saturating_add(replacement.tokens.len().saturating_sub(1)) + }); + let expanded_len = prompt_token_ids.len().saturating_add(replacement_growth); + + let mut expanded = Vec::with_capacity(expanded_len); + let mut ranges = HashMap::>::new(); + + for &token in prompt_token_ids.iter() { + let lane = lanes + .iter_mut() + .find(|lane| lane.marker_token_id == token && !lane.replacements.is_empty()); + let Some(lane) = lane else { + expanded.push(token); + continue; + }; + + let replacement = lane.replacements.pop_front().expect("lane queue is non-empty"); + debug_assert_eq!(replacement.modality, lane.modality); + if replacement.tokens.is_empty() { + bail_multimodal!( + "placeholder token `{}` expanded to no tokens", + lane.placeholder_token + ); + } + + let replacement_len = replacement.tokens.len(); + let is_embed = { + let mask = replacement + .tokens + .iter() + .map(|&token| token as u32 == lane.embed_token_id) + .collect::>(); + WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? + }; + + let expanded_offset = expanded.len(); + expanded.extend(replacement.tokens.iter().map(|&token| token as u32)); + ranges.entry(lane.modality).or_default().push(PlaceholderRange { + offset: expanded_offset, + length: replacement_len, + is_embed: Some(is_embed), + }); + } + + for lane in &lanes { + if !lane.replacements.is_empty() { + bail_multimodal!( + "placeholder token `{}` was not found in tokenized prompt for {} remaining `{}` item(s)", + lane.placeholder_token, + lane.replacements.len(), + lane.modality + ); + } + } + + *prompt_token_ids = expanded; + + Ok(ranges) +} + +#[cfg(test)] +mod tests { + use llm_multimodal::TokenId; + use vllm_engine_core_client::protocol::tensor::WireArrayData; + + use super::super::tests::{ + LLAMA4_IMAGE_END_ID, LLAMA4_IMAGE_ID, LLAMA4_IMAGE_START_ID, LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, LLAMA4_TILE_Y_SEPARATOR_ID, QWEN3_IMAGE_PAD_ID, + QWEN3_VIDEO_PAD_ID, + }; + use super::super::{PreparedMedia, ResolvedPlaceholder}; + use super::*; + + /// Build prepared media directly from placeholder token IDs. + fn prepared_media( + modality: Modality, + placeholder_token: &str, + marker_token_id: u32, + embed_token_id: u32, + replacements: Vec, + ) -> PreparedMedia { + PreparedMedia { + modality, + placeholder: ResolvedPlaceholder { + token: placeholder_token.to_string(), + marker_token_id, + embed_token_id, + }, + replacements, + items: Vec::new(), + } + } + + /// Llama4 image prepared media: the `<|image|>` marker expands to + /// sequences whose embed positions are the `<|patch|>` tokens. + fn llama4_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image|>", + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + replacements, + ) + } + + fn qwen3_image_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + replacements, + ) + } + + fn qwen3_video_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + replacements, + ) + } + + fn llama4_single_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn llama4_multi_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_X_SEPARATOR_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { + let tensor = range.is_embed.as_ref().expect("is_embed mask"); + assert_eq!(tensor.dtype, "bool"); + assert_eq!(tensor.shape, vec![expected.len()]); + assert_eq!( + tensor.data, + WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) + ); + } + + #[test] + fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let prepared = vec![llama4_prepared(vec![llama4_multi_tile_replacement()])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_Y_SEPARATOR_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 8); + assert_bool_mask( + &ranges[0], + &[false, true, false, true, false, false, true, false], + ); + } + + #[test] + fn expand_prompt_tokens_errors_when_placeholder_missing() { + let mut prompt_token_ids = vec![1, 2, 3]; + let prepared = vec![llama4_prepared(vec![llama4_single_tile_replacement()])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + } + + #[test] + fn expand_prompt_tokens_ignores_empty_replacements() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(Vec::new())]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert!(ranges.is_empty()); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_errors_when_replacement_is_empty() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![PromptReplacement::sequence( + Modality::Image, + "<|image|>", + Vec::new(), + )])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!( + matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) + ); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 3, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 5); + assert_bool_mask(&ranges[0], &[false, false, true, true, false]); + assert_eq!(ranges[1].offset, 7); + assert_eq!(ranges[1].length, 5); + assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + } + + #[test] + fn expand_prompt_tokens_interleaves_image_and_video_prepared_media() { + let mut prompt_token_ids = vec![ + 1, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + 4, + ]; + let prepared = vec![ + qwen3_image_prepared(vec![ + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + ), + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 3, + ), + ]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert_eq!( + prompt_token_ids, + vec![ + 1, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 4, + ] + ); + + let image_ranges = &ranges[&Modality::Image]; + assert_eq!(image_ranges[0].offset, 1); + assert_eq!(image_ranges[0].length, 2); + assert_bool_mask(&image_ranges[0], &[true, true]); + assert_eq!(image_ranges[1].offset, 9); + assert_eq!(image_ranges[1].length, 3); + assert_bool_mask(&image_ranges[1], &[true, true, true]); + + let video_ranges = &ranges[&Modality::Video]; + assert_eq!(video_ranges[0].offset, 4); + assert_eq!(video_ranges[0].length, 4); + assert_bool_mask(&video_ranges[0], &[true, true, true, true]); + } + + #[test] + fn expand_prompt_tokens_error_names_modality_with_leftover_replacements() { + let mut prompt_token_ids = vec![1, QWEN3_IMAGE_PAD_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![ + qwen3_image_prepared(vec![PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + )]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("<|video_pad|>") && message.contains("`video`") + )); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } +} diff --git a/rust/src/chat/src/multimodal/image.rs b/rust/src/chat/src/multimodal/image.rs new file mode 100644 index 00000000000..65f86041f73 --- /dev/null +++ b/rust/src/chat/src/multimodal/image.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Image-modality preparation: batch preprocessing and per-item feature +//! build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, ImageFrame, Modality, PreprocessedEncoderInputs}; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +impl MultimodalModelInfo { + /// Preprocess all fetched image frames as one batch and build per-item + /// features. + pub(super) async fn prepare_images( + &self, + frames: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.image.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Image.to_string(), + })?; + let preprocessed = self.preprocess_images(support, &frames).await?; + let replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Image)?; + if replacements.len() != frames.len() { + bail_multimodal!( + "number of image prompt replacements {} does not match number of images {}", + replacements.len(), + frames.len() + ); + } + let items = self.build_image_items(preprocessed, &frames, uuids, model_dtype)?; + + Ok(PreparedMedia { + modality: Modality::Image, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess fetched image frames with the model's resolved vision + /// processor. + /// + /// The processor work is CPU-heavy relative to request wiring, so it runs + /// in a blocking task and returns owned tensors ready for wire + /// conversion. + async fn preprocess_images( + &self, + support: &ModalitySupport, + image_frames: &[Arc], + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); + + // TODO: is it still necessary given that we've already in a dedicated runtime? + tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images, &config)?)) + .await + .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? + } + + /// Convert one batch of preprocessed image tensors into per-item engine + /// kwargs. + /// + /// Tensor fields are sliced per item according to the model spec's field + /// layout declarations. + fn build_image_items( + &self, + preprocessed: PreprocessedEncoderInputs, + frames: &[Arc], + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result> { + let len = frames.len(); + let tensors = tensor::collect_tensors(preprocessed, "pixel_values", model_dtype)?; + + let mut items = Vec::with_capacity(len); + for (index, (frame, uuid)) in izip!(frames, uuids).enumerate() { + let mut data = MmKwargsItem::new(); + for (key, tensor) in &tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); + let (value, field) = match self.spec.field_layouts.get(key) { + Some(FieldLayout::Batched) => ( + tensor.batched_value_at(index)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ), + Some(FieldLayout::Flat { sizes_key }) => { + let sizes = tensors.get(sizes_key).ok_or_else(|| { + multimodal!("flat tensor sizes key `{sizes_key}` is missing") + })?; + let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; + ( + tensor.flat_value_range(start, end)?, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some((end - start) as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } + None => ( + tensor.clone(), + MmField::Shared(MmSharedField { + batch_size: len, + keep_on_cpu, + }), + ), + }; + + data.insert( + key.clone(), + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + items.push(PreparedItem { + data, + hash: frame.hash.clone(), + uuid, + }); + } + + Ok(items) + } +} diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs index eddf8f707e9..f701e0646c6 100644 --- a/rust/src/chat/src/multimodal/tensor.rs +++ b/rust/src/chat/src/multimodal/tensor.rs @@ -1,8 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::collections::HashMap; use half::{bf16, f16}; -use llm_multimodal::{ModelSpecificValue, PreprocessedImages}; -use vllm_engine_core_client::protocol::ModelDtype; +use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; +use vllm_engine_core_client::protocol::dtype::ModelDtype; use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor}; @@ -25,25 +28,31 @@ pub(super) enum KwargValue { Passthrough(ProtocolKwargValue), } -/// Collect `pixel_values` and model-specific outputs into one tensor map. +/// Collect the primary encoder input and model-specific outputs into one +/// tensor map. +/// +/// `primary_key` names the encoder-input tensor as the model's forward kwargs +/// expect it (e.g. `pixel_values` for images, `pixel_values_videos` for +/// videos). pub(super) fn collect_tensors( - preprocessed: PreprocessedImages, + preprocessed: PreprocessedEncoderInputs, + primary_key: &str, float_dtype: ModelDtype, ) -> Result> { - let PreprocessedImages { - pixel_values, + let PreprocessedEncoderInputs { + encoder_input, model_specific, .. } = preprocessed; - let pixel_values = { - let shape = pixel_values.shape().to_vec(); - let data = pixel_values.into_iter().collect(); + let primary_value = { + let shape = encoder_input.shape().to_vec(); + let data = encoder_input.into_iter().collect(); KwargValue::from_f32_tensor(data, shape, float_dtype)? }; let mut tensors = HashMap::new(); - tensors.insert("pixel_values".to_string(), pixel_values); + tensors.insert(primary_key.to_string(), primary_value); for (key, value) in model_specific { tensors.insert(key, KwargValue::from_model_specific(value, float_dtype)?); } @@ -124,10 +133,22 @@ impl TryFrom for ProtocolKwargValue { } impl KwargValue { - /// Extract one image from a batched tensor field. + /// First-axis length for tensor values; `None` for passthrough kwargs. + pub(super) fn first_dim(&self) -> Option { + match self { + Self::F32Tensor { shape, .. } + | Self::F16Tensor { shape, .. } + | Self::Bf16Tensor { shape, .. } + | Self::I64Tensor { shape, .. } + | Self::U32Tensor { shape, .. } => shape.first().copied(), + Self::Passthrough(_) => None, + } + } + + /// Extract one media item from a batched tensor field. /// - /// Batched fields use their first axis as image index and drop that axis in - /// the per-feature value, matching vLLM's batched-field semantics. + /// Batched fields use their first axis as media-item index and drop that + /// axis in the per-feature value, matching vLLM's batched-field semantics. pub(super) fn batched_value_at(&self, index: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -154,9 +175,9 @@ impl KwargValue { } } - /// Extract one image's variable-length range from a flat tensor field. + /// Extract one media item's variable-length range from a flat tensor field. /// - /// Flat fields keep the first axis as the sliced length for this image. + /// Flat fields keep the first axis as the sliced length for this item. pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -184,10 +205,10 @@ impl KwargValue { } } -/// Compute the first-axis range for one image in a flat tensor. +/// Compute the first-axis range for one media item in a flat tensor. /// /// `sizes_key` names a companion tensor whose entries are cumulative slice -/// sizes per image. +/// sizes per media item. pub(super) fn flat_range_for_index( sizes: &KwargValue, sizes_key: &str, @@ -195,7 +216,7 @@ pub(super) fn flat_range_for_index( ) -> Result<(usize, usize)> { let sizes = tensor_as_usize_vec(sizes)?; let size = *sizes.get(index).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` has no entry for image {index}") + multimodal!("flat tensor sizes key `{sizes_key}` has no entry for media item {index}") })?; let start = sizes[..index].iter().sum::(); Ok((start, start + size)) diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs new file mode 100644 index 00000000000..af29799e785 --- /dev/null +++ b/rust/src/chat/src/multimodal/video.rs @@ -0,0 +1,319 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Video-modality preparation: per-clip preprocessing, config resolution, +//! and per-item feature build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, Modality, PreprocessedEncoderInputs, VideoClip}; +use thiserror_ext::AsReport as _; +use tracing::warn; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +/// Forward-kwargs name of the primary video encoder input. +/// +/// Video-capable vLLM models read `pixel_values_videos` alongside +/// `video_grid_thw`, mirroring the HF processor output naming. +const VIDEO_PRIMARY_KEY: &str = "pixel_values_videos"; + +impl MultimodalModelInfo { + /// Preprocess fetched video clips one at a time and build per-item + /// features. + /// + /// Unlike images, each clip runs through the preprocessor independently + /// (a batch of one), so its tensors are complete per item and need no + /// cross-item slicing. + pub(super) async fn prepare_videos( + &self, + clips: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.video.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Video.to_string(), + })?; + let mut replacements = Vec::with_capacity(clips.len()); + let mut items = Vec::with_capacity(clips.len()); + + for (clip, uuid) in izip!(&clips, uuids) { + let preprocessed = self.preprocess_video_clip(support, Arc::clone(clip)).await?; + let mut clip_replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Video)?; + if clip_replacements.len() != 1 { + bail_multimodal!( + "expected exactly one prompt replacement per video clip, got {}", + clip_replacements.len() + ); + } + replacements.push(clip_replacements.pop().unwrap()); + items.push(self.build_video_item( + preprocessed, + clip.hash.clone(), + uuid, + model_dtype, + )?); + } + + Ok(PreparedMedia { + modality: Modality::Video, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess one decoded video clip with the model's resolved vision + /// processor. + async fn preprocess_video_clip( + &self, + support: &ModalitySupport, + clip: Arc, + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + + tokio::task::spawn_blocking(move || { + // Prefer the borrowed-RGB fast path, which avoids materializing a + // `DynamicImage` per sampled frame after media decode. + if let Some(rgb_video) = clip.rgb_video() { + match rgb_video.frame_refs() { + Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs, &config) { + Ok(preprocessed) => return Ok(preprocessed), + Err(error) => warn!( + error = %error.as_report(), + "RGB video preprocessing fast path failed; falling back to materialized frames" + ), + }, + Err(error) => warn!( + error, + "RGB video frame refs are invalid; falling back to materialized frames" + ), + } + } + + let frames = clip.materialized_frames().map_err(|error| multimodal!("{error}"))?; + Ok(processor.preprocess_video(&frames, &config)?) + }) + .await + .map_err(|error| multimodal!("video preprocessing task failed: {error}"))? + } + + /// Convert one preprocessed video clip into engine kwargs. + /// + /// The clip is a batch of one, so no per-item slicing is required: the + /// primary tensor ships as a full-range flat field (the engine re-batches + /// flat fields by concatenating along the declared dim, matching vLLM's + /// `flat_from_sizes` treatment of video patches), and batched metadata + /// tensors drop their singleton batch axis. + fn build_video_item( + &self, + preprocessed: PreprocessedEncoderInputs, + hash: String, + uuid: Option, + model_dtype: ModelDtype, + ) -> Result { + let tensors = tensor::collect_tensors(preprocessed, VIDEO_PRIMARY_KEY, model_dtype)?; + + let mut data = MmKwargsItem::new(); + for (key, tensor) in tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(&key); + let (value, field) = if key == VIDEO_PRIMARY_KEY { + let len = tensor + .first_dim() + .ok_or_else(|| multimodal!("video encoder input `{key}` is not a tensor"))?; + ( + tensor, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(len as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } else if matches!( + self.spec.field_layouts.get(&key), + Some(FieldLayout::Batched) + ) { + ( + tensor.batched_value_at(0)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ) + } else { + ( + tensor, + MmField::Shared(MmSharedField { + batch_size: 1, + keep_on_cpu, + }), + ) + }; + + data.insert( + key, + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + Ok(PreparedItem { data, hash, uuid }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use llm_multimodal::ModelSpecificValue; + use ndarray::ArrayD; + use vllm_engine_core_client::protocol::multimodal::MmKwargValue; + + use super::super::tests::{ + QWEN3_IMAGE_PAD_ID, QWEN3_VIDEO_PAD_ID, qwen3_vl_info, qwen3_vl_tokenizer, + }; + use super::super::{MultimodalConfigFiles, MultimodalModelInfo}; + use super::*; + + #[test] + fn from_paths_resolves_video_config_from_dedicated_file_or_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + }) + .to_string(), + ) + .unwrap(); + + let info_for = |files: MultimodalConfigFiles<'_>| { + MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + files, + Arc::new(qwen3_vl_tokenizer()), + ) + }; + + // Dedicated video preprocessor config file. + let video_config_path = dir.path().join("video_preprocessor_config.json"); + std::fs::write(&video_config_path, r#"{"size":{"shortest_edge":128}}"#).unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // `video_processor` section of the combined processor config. + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"video_processor":{"size":{"shortest_edge":128}}}"#, + ) + .unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Neither source: video support still resolves on the image config. + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Malformed dedicated file is a real error, not a silent fallback. + std::fs::write(&video_config_path, r#"{"size""#).unwrap(); + let error = match info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) { + Err(error) => error, + Ok(_) => panic!("malformed video preprocessor config should fail"), + }; + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("failed to parse video_preprocessor_config.json") + )); + } + + #[test] + fn build_video_item_names_primary_tensor_and_layouts() { + let info = qwen3_vl_info(); + // One clip flattened to 6 patches with 4 features each. + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::zeros(vec![6, 4]), + feature_token_counts: vec![6], + item_sizes: vec![(32, 32)], + model_specific: HashMap::from([ + ( + "video_grid_thw".to_string(), + ModelSpecificValue::int_2d(vec![1, 2, 3], 1, 3), + ), + ( + "patches_per_video".to_string(), + ModelSpecificValue::int_1d(vec![6]), + ), + ]), + }; + + let item = info + .build_video_item( + preprocessed, + "".to_string(), + None, + ModelDtype::Float32, + ) + .unwrap(); + + let primary = &item.data[VIDEO_PRIMARY_KEY]; + assert!(matches!( + &primary.field, + MmField::Flat(MmFlatField { slices, dim: 0, .. }) + if matches!( + slices.as_slice(), + [MmSlice::Slice(SliceSpec { start: Some(0), stop: Some(6), step: None })] + ) + )); + + // Batched metadata drops its singleton batch axis per item. + let grid = &item.data["video_grid_thw"]; + assert!(matches!(&grid.field, MmField::Batched(_))); + let MmKwargValue::Tensor(grid_tensor) = grid.data.as_ref().unwrap() else { + panic!("expected tensor value for video_grid_thw"); + }; + assert_eq!(grid_tensor.shape, vec![3]); + + assert_eq!(item.hash, ""); + } +} diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 40526a9e84c..ab288336f4c 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -1,42 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Default output processing pipeline. -mod reasoning; -mod tool; +mod structural_tag; +mod unified; use std::sync::Once; -use futures::{Stream, StreamExt as _}; +use futures::StreamExt as _; use tracing::info; -use trait_set::trait_set; +use vllm_parser::unified::{CombinedParser, UnifiedParser}; use vllm_text::tokenizer::DynTokenizer; -use self::reasoning::reasoning_event_stream; -use self::tool::tool_event_stream; +use self::structural_tag::apply_structural_tag_constraint; +use self::unified::unified_event_stream; use super::structured::structured_chat_event_stream; use crate::error::Result; -use crate::output::{ - AssistantEvent, ChatOutputProcessor, ContentEvent, DynChatEventStream, - DynDecodedTextEventStream, -}; +use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream}; use crate::parser::ParserSelection; use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; -use crate::request::{ChatRequest, ChatToolChoice}; +use crate::parser::unified::UnifiedParserFactory; +use crate::request::{ChatRequest, ChatTool}; use crate::{Error, Result as ChatResult}; -trait_set! { - trait ContentEventStream = Stream> + Send + 'static; -} - /// Default request-scoped output processor used by Hugging Face style chat /// backends. /// /// This implementation assumes the backend already emitted decoded text deltas, -/// then optionally layers reasoning parsing and tool-call parsing before +/// then optionally layers unified reasoning and tool-call parsing before /// assembling final structured chat events. pub struct DefaultChatOutputProcessor { - reasoning_parser: Option>, - tool_parser: Option>, + parser: Box, + parallel_tool_calls: bool, } impl DefaultChatOutputProcessor { @@ -53,27 +50,39 @@ impl DefaultChatOutputProcessor { tool_call_parser: &ParserSelection, reasoning_parser: &ParserSelection, ) -> ChatResult { - let tool_parsing_enabled = - matches!(request.tool_choice, ChatToolChoice::Auto) && !request.tools.is_empty(); - let tool_parser = if tool_parsing_enabled { - Some(Self::resolve_tool_parser( - request, + let parser = if tool_call_parser == reasoning_parser + && let Some(parser) = Self::resolve_optional_unified_parser( + &request.tools, model_id, + tokenizer.clone(), tool_call_parser, - )?) + )? { + parser } else { - None + let tool_parsing_enabled = request.tool_parsing_enabled(); + let tool_parser = if tool_parsing_enabled { + Some(Self::resolve_tool_parser( + &request.tools, + model_id, + tool_call_parser, + )?) + } else { + None + }; + let reasoning_parser = + Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; + Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box }; - let reasoning_parser = Self::resolve_optional_reasoning_parser( - request, - model_id, - tokenizer, - reasoning_parser, - )?; + + apply_structural_tag_constraint(request, parser.structural_tag_model())?; + + if parser.preserve_special_tokens() { + request.decode_options.skip_special_tokens = false; + } Ok(Self { - reasoning_parser, - tool_parser, + parser, + parallel_tool_calls: request.parallel_tool_calls, }) } @@ -84,13 +93,13 @@ impl DefaultChatOutputProcessor { /// content is treated as opaque text. pub fn plain_text_only() -> Self { Self { - reasoning_parser: None, - tool_parser: None, + parser: Box::new(CombinedParser::plain_text_only()), + parallel_tool_calls: true, } } fn resolve_tool_parser( - request: &mut ChatRequest, + tools: &[ChatTool], model_id: &str, selection: &ParserSelection, ) -> ChatResult> { @@ -106,18 +115,37 @@ impl DefaultChatOutputProcessor { ParserSelection::Explicit(name) => name.as_str(), }; - let parser = factory.create(parser_name, &request.tools)?; - - if parser.preserve_special_tokens() { - request.decode_options.skip_special_tokens = false; - } + let parser = factory.create(parser_name, tools)?; TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser")); Ok(parser) } + fn resolve_optional_unified_parser( + tools: &[ChatTool], + model_id: &str, + tokenizer: DynTokenizer, + selection: &ParserSelection, + ) -> ChatResult>> { + let factory = UnifiedParserFactory::global(); + let parser_name = match selection { + ParserSelection::Auto => factory.resolve_name_for_model(model_id), + ParserSelection::None => None, + ParserSelection::Explicit(name) if factory.contains(name) => Some(name.as_str()), + ParserSelection::Explicit(_) => None, + }; + + let Some(parser_name) = parser_name else { + return Ok(None); + }; + + let parser = factory.create(parser_name, tools, tokenizer)?; + + UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser")); + Ok(Some(parser)) + } + fn resolve_optional_reasoning_parser( - request: &mut ChatRequest, model_id: &str, tokenizer: DynTokenizer, selection: &ParserSelection, @@ -136,10 +164,6 @@ impl DefaultChatOutputProcessor { let parser = factory.create(parser_name, tokenizer)?; - if parser.preserve_special_tokens() { - request.decode_options.skip_special_tokens = false; - } - REASONING_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using reasoning parser")); Ok(Some(parser)) } @@ -147,20 +171,91 @@ impl DefaultChatOutputProcessor { static TOOL_PARSER_LOG_ONCE: Once = Once::new(); static REASONING_PARSER_LOG_ONCE: Once = Once::new(); +static UNIFIED_PARSER_LOG_ONCE: Once = Once::new(); impl ChatOutputProcessor for DefaultChatOutputProcessor { /// Transforms a raw generate-output token stream into structured chat - /// events through three sequential stages once text decoding has + /// events through two sequential stages once text decoding has /// already happened: /// - /// 1. [`reasoning_event_stream`] — reasoning/content separation - /// 2. [`tool_event_stream`] — tool-call parsing - /// 3. [`structured_chat_event_stream`] — final block assembly + /// 1. [`unified_event_stream`] — reasoning and tool-call parsing + /// 2. [`structured_chat_event_stream`] — final block assembly fn process(self: Box, decoded: DynDecodedTextEventStream) -> Result { - let reasoning = reasoning_event_stream(decoded, self.reasoning_parser); - let tool = tool_event_stream(reasoning, self.tool_parser); - let structured = structured_chat_event_stream(tool); + let parsed = unified_event_stream(decoded, self.parser); + let structured = structured_chat_event_stream(parsed, self.parallel_tool_calls); Ok(structured.boxed()) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::DefaultChatOutputProcessor; + use crate::Error; + use crate::parser::ParserSelection; + use crate::request::ChatRequest; + + fn tokenizer() -> Arc { + Arc::new( + TestTokenizer::new() + .with_regular_token("<|channel>", 256) + .with_regular_token("", 257), + ) + } + + #[test] + fn equal_explicit_gemma4_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + let selection = ParserSelection::Explicit("gemma4".to_string()); + + DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &selection, + &selection, + ) + .unwrap(); + } + + #[test] + fn auto_auto_gemma4_model_uses_unified_parser() { + let mut request = ChatRequest::for_test(); + + DefaultChatOutputProcessor::new( + &mut request, + "google/gemma-4-27b-it", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Auto, + ) + .unwrap(); + } + + #[test] + fn mixed_gemma4_selection_uses_split_dummy_error() { + let mut request = ChatRequest::for_test(); + let error = match DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &ParserSelection::Auto, + &ParserSelection::Explicit("gemma4".to_string()), + ) { + Ok(_) => panic!("expected mixed Gemma4 parser selection to fail"), + Err(error) => error, + }; + + let Error::ParserInitialization { error, .. } = error else { + panic!("expected parser initialization error"); + }; + assert_eq!( + error.to_string(), + "`gemma4` only provides a unified parser; the same reasoning parser and tool parser should be specified together" + ); + } +} diff --git a/rust/src/chat/src/output/default/reasoning.rs b/rust/src/chat/src/output/default/reasoning.rs deleted file mode 100644 index b51ce41961d..00000000000 --- a/rust/src/chat/src/output/default/reasoning.rs +++ /dev/null @@ -1,504 +0,0 @@ -//! Adapts decoded text updates into reasoning-aware assistant deltas. -//! -//! This stage sits between low-level token decoding and final block assembly. -//! It is the only place in the new pipeline that understands reasoning -//! separation: `decoded.rs` still only produces plain text deltas, while later -//! stages consume the semantic `Text` / `Reasoning` split emitted here. - -use asynk_strim_attr::{TryYielder, try_stream}; -use futures::{StreamExt as _, pin_mut}; -use thiserror_ext::AsReport; -use tracing::warn; -use vllm_text::output::DecodedTextEvent; - -use super::ContentEvent; -use crate::Result; -use crate::error::Error; -use crate::event::AssistantBlockKind; -use crate::output::DecodedTextEventStream; -use crate::parser::reasoning::{ReasoningDelta, ReasoningParser}; - -/// Per-stream reasoning parsing state. -struct ReasoningState { - /// Reasoning parser for the current model family. - parser: Box, - /// Whether reasoning parsing has already failed for this stream. - parser_failed: bool, -} - -impl ReasoningState { - /// Create one fresh reasoning-adaptation state for a new streamed response. - fn new(parser: Box) -> Self { - Self { - parser, - parser_failed: false, - } - } - - /// Convert one decoded text delta into zero or more semantic assistant - /// deltas. - fn process_delta(&mut self, delta: String) -> Vec { - // If the parser has already failed, skip parsing and return plain text deltas. - if self.parser_failed { - return vec![ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }]; - } - - let mut events = Vec::new(); - - match self.parser.push(&delta) { - Ok(result) => { - push_reasoning_delta(&mut events, result); - } - Err(error) => { - if !self.parser_failed { - warn!( - error = %error.as_report(), - "reasoning parser failed; falling back to plain text deltas" - ); - self.parser_failed = true; - } - push_text_delta(&mut events, AssistantBlockKind::Text, delta); - } - } - - events - } - - /// Initialize parser state once prompt token IDs are available. - fn initialize(&mut self, prompt_token_ids: &[u32]) { - if self.parser_failed { - return; - } - - match self.parser.initialize(prompt_token_ids) { - Ok(()) => {} - Err(error) => { - warn!( - error = %error.as_report(), - "failed to initialize reasoning parser; falling back to plain text deltas" - ); - self.parser_failed = true; - } - } - } - - /// Flush any parser-held partial delimiter state at end of stream. - fn finish(&mut self) -> Vec { - if self.parser_failed { - return Vec::new(); - } - - match self.parser.finish() { - Ok(result) => { - let mut events = Vec::new(); - push_reasoning_delta(&mut events, result); - events - } - Err(error) => { - warn!(error = %error.as_report(), "failed to flush reasoning parser state"); - Vec::new() - } - } - } -} - -/// Push one semantic text delta if it is non-empty. -fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { - if delta.is_empty() { - return; - } - events.push(ContentEvent::TextDelta { kind, delta }); -} - -/// Convert one parsed reasoning delta into zero or more content events. -fn push_reasoning_delta(events: &mut Vec, delta: ReasoningDelta) { - if let Some(reasoning) = delta.reasoning { - push_text_delta(events, AssistantBlockKind::Reasoning, reasoning); - } - if let Some(content) = delta.content { - push_text_delta(events, AssistantBlockKind::Text, content); - } -} - -/// Wrap one decoded-text stream into the internal reasoning event stream. -#[try_stream] -pub(crate) async fn reasoning_event_stream( - decoded_stream: impl DecodedTextEventStream, - reasoning_parser: Option>, - mut y: TryYielder, -) -> Result<()> { - pin_mut!(decoded_stream); - - // Without a parser, pass through as plain text deltas. - let Some(reasoning_parser) = reasoning_parser else { - while let Some(event) = decoded_stream.next().await.transpose()? { - for next in ContentEvent::from_decoded_plain_text(event) { - y.yield_ok(next).await; - } - } - return Ok(()); - }; - - let mut state = ReasoningState::new(reasoning_parser); - - while let Some(event) = decoded_stream.next().await.transpose()? { - match event { - DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => { - state.initialize(&prompt_token_ids); - y.yield_ok(ContentEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) - .await; - } - DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - } => { - for next in state.process_delta(delta) { - y.yield_ok(next).await; - } - if logprobs.is_some() || !token_ids.is_empty() { - y.yield_ok(ContentEvent::LogprobsDelta { - logprobs, - token_ids, - }) - .await; - } - if let Some(finished) = finished { - for next in state.finish() { - y.yield_ok(next).await; - } - y.yield_ok(ContentEvent::Done { - prompt_token_count: finished.prompt_token_count, - output_token_count: finished.output_token_count, - finish_reason: finished.finish_reason, - kv_transfer_params: finished.kv_transfer_params, - }) - .await; - } - } - } - } - Ok(()) -} - -#[cfg(test)] -mod tests { - - use std::sync::Arc; - - use futures::{StreamExt as _, stream}; - use vllm_llm::FinishReason; - use vllm_text::output::{ - DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, DecodedTokenLogprob, - }; - use vllm_tokenizer::{DynTokenizer, Tokenizer}; - - use super::super::ContentEvent; - use super::reasoning_event_stream; - use crate::event::AssistantBlockKind; - use crate::parser::reasoning::{ - ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory, names, - }; - - struct FakeTokenizer; - - impl Tokenizer for FakeTokenizer { - fn encode( - &self, - text: &str, - _add_special_tokens: bool, - ) -> vllm_tokenizer::Result> { - Ok(text.chars().map(u32::from).collect()) - } - - fn decode( - &self, - token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_tokenizer::Result { - Ok(token_ids - .iter() - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) - .collect()) - } - - fn token_to_id(&self, token: &str) -> Option { - match token { - "" => Some(1), - "" => Some(2), - _ => None, - } - } - } - - struct FailingReasoningParser { - fail_next: bool, - } - - impl ReasoningParser for FailingReasoningParser { - fn create(_tokenizer: DynTokenizer) -> Result, ReasoningError> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { fail_next: true })) - } - - fn push(&mut self, _text: &str) -> Result { - if self.fail_next { - self.fail_next = false; - return Err(ReasoningError::MissingToken { - token: "".to_string(), - }); - } - Ok(ReasoningDelta::default()) - } - } - - fn test_reasoning_parser(factory: &mut ReasoningParserFactory) -> Box { - factory.register_parser::("failing"); - - factory.create("failing", Arc::new(FakeTokenizer)).unwrap() - } - - #[tokio::test] - async fn reasoning_parser_failure_falls_back_to_plain_text() { - let mut factory = ReasoningParserFactory::new(); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "abc".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "def".to_string(), - token_ids: vec![], - logprobs: None, - finished: Some(vllm_text::Finished { - prompt_token_count: 3, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - }), - ]); - - let collected = reasoning_event_stream(events, Some(test_reasoning_parser(&mut factory))) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }, - ContentEvent::Done { - prompt_token_count: 3, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - } - - #[tokio::test] - async fn reasoning_stream_preserves_logprobs_delta() { - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "abc".to_string(), - token_ids: vec![], - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.1, - rank: 1, - }], - }], - }), - finished: None, - }), - ]); - - let collected = reasoning_event_stream(events, None) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert_eq!( - collected, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - ContentEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.1, - rank: 1, - }], - }], - }), - token_ids: vec![], - }, - ] - ); - } - - #[tokio::test] - async fn qwen3_parser_uses_prompt_end_marker_to_switch_to_content() { - let tokenizer = Arc::new(FakeTokenizer); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![2].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "thought ".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "doneOK".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - ]); - - let factory = ReasoningParserFactory::new(); - let collected = reasoning_event_stream( - events, - Some(factory.create(names::QWEN3, tokenizer).unwrap()), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![2].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "thought ".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "doneOK".to_string(), - }, - ] - ); - } - - #[tokio::test] - async fn qwen3_parser_tolerates_prompt_prefill_reasoning() { - let tokenizer = Arc::new(FakeTokenizer); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "thought ".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "doneOK".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - ]); - - let factory = ReasoningParserFactory::new(); - let collected = reasoning_event_stream( - events, - Some(factory.create(names::QWEN3, tokenizer).unwrap()), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Reasoning, - delta: "thought ".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Reasoning, - delta: "done".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "OK".to_string(), - }, - ] - ); - } -} diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs new file mode 100644 index 00000000000..79aa1a1f1b5 --- /dev/null +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Applies xgrammar structural-tag constraints for strict tool calling. + +use thiserror_ext::AsReport; +use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputBackend, StructuredOutputsParams, +}; +use vllm_parser::tool::StructuralTagModel; +use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, + build_structural_tag, +}; + +use crate::request::{ChatRequest, ChatToolChoice}; +use crate::{Error, Result as ChatResult}; + +/// Apply structural tag constraints to the request based on the tool parser's structural tag +/// support and the request's tool choice. +pub(super) fn apply_structural_tag_constraint( + request: &mut ChatRequest, + model: Option, +) -> ChatResult<()> { + let Some(model) = model else { + return Ok(()); + }; + let Some(tool_choice) = structural_tag_tool_choice(request) else { + return Ok(()); + }; + + let tools = request + .tools + .iter() + .map(|tool| { + ToolParam::Function(FunctionToolParam::new(FunctionDefinition { + name: tool.name.clone(), + description: tool.description.clone(), + parameters: Some(tool.parameters.clone()), + strict: tool.strict, + })) + }) + .collect::>(); + + let structural_tag = build_structural_tag(model, &tools, tool_choice, false) + .and_then(|tag| tag.to_json_string()) + .map_err(|error| Error::StructuralTag { + message: error.to_report_string(), + })?; + + // Overwrite any existing structured output settings with the structural tag constraint. + request.sampling_params.structured_outputs = Some(StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..StructuredOutputsParams::structural_tag(structural_tag) + }); + + Ok(()) +} + +/// Resolve the tool choice used for [`xgrammar_structural_tag`] based on the request. +/// +/// Returns `None` if no structural tag constraints should be applied. +fn structural_tag_tool_choice(request: &ChatRequest) -> Option { + if request.tools.is_empty() { + return None; + } + + match &request.tool_choice { + // For `Auto`, only apply the structural tag if there's at least one strict tool. + ChatToolChoice::Auto if request.tools.iter().any(|tool| tool.strict == Some(true)) => { + Some(StructuralTagToolChoice::auto()) + } + ChatToolChoice::Auto | ChatToolChoice::None => None, + + ChatToolChoice::Required => Some(StructuralTagToolChoice::required()), + ChatToolChoice::Function { name } => Some(StructuralTagToolChoice::function(name.clone())), + } +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputBackend, StructuredOutputsParams, + }; + use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser}; + + use super::*; + + fn chat_tool(name: &str, strict: Option) -> Tool { + Tool { + name: name.to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + }), + strict, + } + } + + fn qwen3_coder_parser(tools: &[Tool]) -> Box { + Qwen3CoderToolParser::create(tools).expect("Qwen3 Coder parser should build") + } + + fn request(tool_choice: ChatToolChoice, tools: Vec) -> ChatRequest { + ChatRequest { + tool_choice, + tools, + ..ChatRequest::for_test() + } + } + + fn structural_tag_value(request: &ChatRequest) -> Value { + let params = request + .sampling_params + .structured_outputs + .as_ref() + .expect("structured outputs should be set"); + assert_eq!(params.backend, StructuredOutputBackend::Xgrammar); + let structural_tag = params + .constraint + .as_structural_tag() + .expect("structured output constraint should be structural_tag"); + serde_json::from_str(structural_tag).expect("structural_tag should be valid JSON") + } + + fn structured_outputs(request: &ChatRequest) -> &StructuredOutputsParams { + request + .sampling_params + .structured_outputs + .as_ref() + .expect("structured outputs should be set") + } + + #[test] + fn auto_strict_tool_choice_builds_structural_tag() { + let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag should build"); + + let tag = structural_tag_value(&request); + assert_eq!(tag["type"], "structural_tag"); + assert!(tag.to_string().contains("search")); + } + + #[test] + fn auto_non_strict_tool_choice_skips_structural_tag() { + let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag decision should succeed"); + + assert!(request.sampling_params.structured_outputs.is_none()); + } + + #[test] + fn auto_strict_tool_choice_overwrites_existing_json_guidance() { + let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); + request.sampling_params.structured_outputs = Some(StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..StructuredOutputsParams::json(json!({"type": "object"})) + }); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag should build"); + + let params = structured_outputs(&request); + assert!(params.constraint.is_structural_tag()); + let tag = structural_tag_value(&request); + assert_eq!(tag["type"], "structural_tag"); + assert!(tag.to_string().contains("search")); + } + + #[test] + fn required_tool_choice_builds_structural_tag_without_strict_tools() { + let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag should build"); + + let tag = structural_tag_value(&request); + assert_eq!(tag["type"], "structural_tag"); + assert!(tag.to_string().contains("search")); + } + + #[test] + fn required_tool_choice_overwrites_existing_json_object_guidance() { + let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); + request.sampling_params.structured_outputs = Some(StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..StructuredOutputsParams::json_object() + }); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag should build"); + + let params = structured_outputs(&request); + assert!(params.constraint.is_structural_tag()); + let tag = structural_tag_value(&request); + assert_eq!(tag["type"], "structural_tag"); + assert!(tag.to_string().contains("search")); + } + + #[test] + fn named_tool_choice_builds_structural_tag_for_named_tool_only() { + let mut request = request( + ChatToolChoice::Function { + name: "lookup".to_string(), + }, + vec![chat_tool("search", None), chat_tool("lookup", None)], + ); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag should build"); + + let tag = structural_tag_value(&request).to_string(); + assert!(tag.contains("lookup")); + assert!(!tag.contains("search")); + } + + #[test] + fn none_tool_choice_skips_structural_tag() { + let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag decision should succeed"); + + assert!(request.sampling_params.structured_outputs.is_none()); + } + + #[test] + fn none_tool_choice_preserves_existing_json_object_guidance() { + let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); + request.sampling_params.structured_outputs = Some(StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..StructuredOutputsParams::json_object() + }); + let parser = qwen3_coder_parser(&request.tools); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) + .expect("structural tag decision should succeed"); + + let params = structured_outputs(&request); + assert!(params.constraint.is_json_object()); + } +} diff --git a/rust/src/chat/src/output/default/tool.rs b/rust/src/chat/src/output/default/tool.rs deleted file mode 100644 index 89696675306..00000000000 --- a/rust/src/chat/src/output/default/tool.rs +++ /dev/null @@ -1,961 +0,0 @@ -//! Adapts plain assistant text deltas into tool-call-aware assistant updates. -//! -//! This stage runs after reasoning separation and before final block assembly. -//! It only inspects normal assistant text, leaves reasoning deltas untouched, -//! and translates incremental tool parsing output into internal tool-call -//! events while preserving plain-text fallback behavior. - -use asynk_strim_attr::{TryYielder, try_stream}; -use futures::{StreamExt as _, pin_mut}; -use thiserror_ext::AsReport; -use tracing::warn; - -use super::{AssistantEvent, ContentEvent, ContentEventStream}; -use crate::Result; -use crate::error::Error; -use crate::event::AssistantBlockKind; -use crate::output::generate_tool_call_id; -use crate::parser::tool::{ToolCallDelta, ToolParser, ToolParserOutput}; - -/// Per-stream tool parsing state. -struct ToolState { - /// Parser for the current model family. - parser: Box, - /// Whether tool parsing has already failed for this stream. - parser_failed: bool, - /// The parser-local index of the currently open tool call, if any. - // NOTE: We only allow single open tool call at a time right now, since that's what all - // supported parsers currently emit. Change this to a `BTreeMap` if we need to support multiple - // interleaved calls in the future. - open_call_index: Option, -} - -impl ToolState { - /// Create one fresh tool-parsing state for a new streamed response. - fn new(parser: Box) -> Self { - Self { - parser, - parser_failed: false, - open_call_index: None, - } - } - - /// Convert one semantic assistant text delta into zero or more tool-aware - /// internal events. - fn process_text_delta( - &mut self, - kind: AssistantBlockKind, - delta: String, - ) -> Result> { - let mut events = Vec::new(); - - // Only normal assistant text is eligible for tool parsing. Reasoning - // blocks and plain-text fallback should pass through unchanged. - if kind != AssistantBlockKind::Text || self.parser_failed { - self.open_call_index = None; - events.push(AssistantEvent::TextDelta { kind, delta }); - return Ok(events); - } - - let mut output = ToolParserOutput::default(); - let parse_result = self.parser.parse_into(&delta, &mut output); - - match parse_result { - Ok(()) => self.process_parser_output(kind, output, &mut events)?, - Err(error) => { - warn!( - error = %error.as_report(), - "tool parser failed; falling back to plain text deltas" - ); - // Permanently mark this parser as failed. - // TODO: we may consider recovering from parsing errors in the future. - self.parser_failed = true; - - // On parsing failure, we still apply the partial parser output if any, but we close - // any open tool calls and emit the remaining buffered text as a plain-text delta to - // preserve as much of the output as possible. - self.process_parser_output(kind, output, &mut events)?; - self.open_call_index = None; - push_text_delta(&mut events, kind, self.parser.reset()); - } - } - - Ok(events) - } - - /// Apply one parsed tool output to the current stream state. - fn process_parser_output( - &mut self, - kind: AssistantBlockKind, - output: ToolParserOutput, - events: &mut Vec, - ) -> Result<()> { - // When we are not currently streaming a tool call, preserve plain - // text first and then surface any new tool call items. - if self.open_call_index.is_none() { - push_text_delta(events, kind, output.normal_text); - self.process_tool_items(output.calls, events)?; - } else { - // Once a tool call is open, prioritize tool deltas first. If the - // parser emits normal text again, close the tool call and resume - // plain text output. - self.process_tool_items(output.calls, events)?; - if !output.normal_text.is_empty() { - self.open_call_index = None; - push_text_delta(events, kind, output.normal_text); - } - } - Ok(()) - } - - /// Apply one batch of parsed tool-call deltas emitted by the parser. - fn process_tool_items( - &mut self, - items: Vec, - events: &mut Vec, - ) -> Result<()> { - for item in items { - if let Some(name) = item.name { - let is_new_tool = match self.open_call_index { - Some(open_call_index) => open_call_index != item.tool_index, - None => true, - }; - if is_new_tool { - let id = self - .parser - .tool_call_id(item.tool_index) - .map(str::to_string) - .unwrap_or_else(generate_tool_call_id); - self.open_call_index = Some(item.tool_index); - events.push(AssistantEvent::ToolCallStart { id, name }); - } - } - - if item.arguments.is_empty() { - // No arguments delta to apply. - continue; - } - let Some(open_call_index) = self.open_call_index else { - return Err(Error::ToolCallStreamInvariant { - message: format!( - "received arguments for tool index {} before any tool-call start", - item.tool_index - ), - }); - }; - if open_call_index != item.tool_index { - return Err(Error::ToolCallStreamInvariant { - message: format!( - "received arguments for tool index {} while tool index {} is open", - item.tool_index, open_call_index - ), - }); - } - - events.push(AssistantEvent::ToolCallArgumentsDelta { - delta: item.arguments, - }); - } - Ok(()) - } - - /// Flush parser state at end-of-stream and close any remaining open calls. - fn finish(&mut self) -> Result> { - let mut events = Vec::new(); - - if self.parser_failed { - return Ok(events); - } - - match self.parser.finish() { - Ok(output) => { - self.process_parser_output(AssistantBlockKind::Text, output, &mut events)? - } - Err(error) => { - warn!( - error = %error.as_report(), - "tool parser finish failed; closing open tool calls with buffered state" - ); - self.parser_failed = true; - } - } - - Ok(events) - } -} - -/// Push one plain-text delta if it is non-empty. -fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { - if delta.is_empty() { - return; - } - events.push(AssistantEvent::TextDelta { kind, delta }); -} - -/// Wrap one semantic assistant stream into the internal tool-aware assistant -/// stream. -#[try_stream] -pub(crate) async fn tool_event_stream( - stream: impl ContentEventStream, - parser: Option>, - mut y: TryYielder, -) -> Result<()> { - // Without a parser, pass through the input stream unchanged. - let Some(parser) = parser else { - pin_mut!(stream); - while let Some(event) = stream.next().await.transpose()? { - y.yield_ok(event.into()).await; - } - return Ok(()); - }; - - pin_mut!(stream); - let mut state = ToolState::new(parser); - - while let Some(event) = stream.next().await.transpose()? { - match event { - ContentEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => { - y.yield_ok(AssistantEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) - .await; - } - ContentEvent::TextDelta { kind, delta } => { - for next in state.process_text_delta(kind, delta)? { - y.yield_ok(next).await; - } - } - ContentEvent::LogprobsDelta { - logprobs, - token_ids, - } => { - y.yield_ok(AssistantEvent::LogprobsDelta { - logprobs, - token_ids, - }) - .await; - } - ContentEvent::Done { - prompt_token_count, - output_token_count, - finish_reason, - kv_transfer_params, - } => { - for next in state.finish()? { - y.yield_ok(next).await; - } - - y.yield_ok(AssistantEvent::Done { - prompt_token_count, - output_token_count, - finish_reason, - kv_transfer_params, - }) - .await; - } - } - } - Ok(()) -} - -#[cfg(test)] -mod tests { - - use futures::{StreamExt as _, stream}; - use vllm_llm::FinishReason; - use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; - use vllm_tool_parser::Result; - - use super::super::{AssistantEvent, ContentEvent}; - use super::tool_event_stream; - use crate::error::Error; - use crate::event::{AssistantBlockKind, AssistantMessageExt as _}; - use crate::output::structured::structured_chat_event_stream; - use crate::parser::tool::{ - DeepSeekV4ToolParser, ToolParser, ToolParserError, ToolParserOutput, - }; - use crate::request::ChatTool; - use crate::stream::{ChatEventStream, CollectedAssistantMessage}; - - struct FailingParser { - fail_next: bool, - buffered: String, - } - - struct ScriptedParser { - push_outputs: Vec, - finish_output: ToolParserOutput, - } - - struct PartialThenFailParser { - buffered: String, - } - - struct IdScriptedParser { - output: ToolParserOutput, - tool_call_id: Option, - } - - impl ToolParser for FailingParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - fail_next: false, - buffered: String::new(), - })) - } - - fn parse_into(&mut self, chunk: &str, _output: &mut ToolParserOutput) -> Result<()> { - self.buffered.push_str(chunk); - if self.fail_next { - self.fail_next = false; - return Err(ToolParserError::ParsingFailed { - message: "boom".to_string(), - }); - } - - self.buffered.clear(); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - std::mem::take(&mut self.buffered) - } - } - - impl ToolParser for ScriptedParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - push_outputs: Vec::new(), - finish_output: ToolParserOutput::default(), - })) - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - let mut next = self.push_outputs.pop().unwrap_or_default(); - output.normal_text.push_str(&next.normal_text); - output.calls.append(&mut next.calls); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(std::mem::take(&mut self.finish_output)) - } - - fn reset(&mut self) -> String { - String::new() - } - } - - impl ToolParser for IdScriptedParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - output: ToolParserOutput::default(), - tool_call_id: None, - })) - } - - fn tool_call_id(&self, tool_index: usize) -> Option<&str> { - (tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten() - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - output.append(std::mem::take(&mut self.output)); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - String::new() - } - } - - impl ToolParser for PartialThenFailParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - buffered: String::new(), - })) - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - output.calls.extend([ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: String::new(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: None, - arguments: r#"{"location":"SF"}"#.to_string(), - }, - ]); - self.buffered.push_str(" trailing text"); - Err(ToolParserError::ParsingFailed { - message: "boom".to_string(), - }) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - std::mem::take(&mut self.buffered) - } - } - - fn deepseek_v4_test_tools() -> Vec { - vec![ - ChatTool { - name: "get_weather".to_string(), - description: None, - parameters: serde_json::json!({ - "type": "object", - "properties": { - "location": { "type": "string" } - } - }), - strict: None, - }, - ChatTool { - name: "add".to_string(), - description: None, - parameters: serde_json::json!({ - "type": "object", - "properties": { - "x": { "type": "integer" }, - "y": { "type": "integer" } - } - }), - strict: None, - }, - ] - } - - async fn collect_deepseek_v4_message(chunks: Vec) -> CollectedAssistantMessage { - let events = chunks - .into_iter() - .map(|delta| { - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }) - }) - .chain(std::iter::once(Ok(ContentEvent::Done { - prompt_token_count: 1, - output_token_count: 1, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }))); - let parser = DeepSeekV4ToolParser::create(&deepseek_v4_test_tools()).unwrap(); - let assistant_events = tool_event_stream(stream::iter(events), Some(parser)); - let chat_events = structured_chat_event_stream(assistant_events); - - ChatEventStream::new("req_deepseek_v4".to_string(), Box::pin(chat_events)) - .collect_message() - .await - .unwrap() - } - - fn message_tool_projection( - message: &CollectedAssistantMessage, - ) -> (String, Vec<(String, serde_json::Value)>) { - ( - message.message.text(), - message - .message - .tool_calls() - .map(|call| { - ( - call.name.clone(), - serde_json::from_str(&call.arguments).unwrap(), - ) - }) - .collect(), - ) - } - - #[tokio::test] - async fn tool_parser_error_preserves_partial_output_and_flushes_buffer() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - prompt_token_count: 1, - output_token_count: 1, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let events = tool_event_stream( - events, - Some(Box::new(PartialThenFailParser { - buffered: String::new(), - })), - ) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { name, .. } if name == "get_weather" - )); - assert!(matches!( - &events[1], - AssistantEvent::ToolCallArgumentsDelta { delta } if delta == r#"{"location":"SF"}"# - )); - assert_eq!( - events[2], - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: " trailing text".to_string(), - } - ); - assert!(matches!(events[3], AssistantEvent::Done { .. })); - } - - #[tokio::test] - async fn tool_stream_preserves_parser_provided_tool_call_id() { - let events = stream::iter(vec![Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - })]); - let parser = IdScriptedParser { - output: ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: "{}".to_string(), - }], - }, - tool_call_id: Some("functions.get_weather:0".to_string()), - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { id, name } - if id == "functions.get_weather:0" && name == "get_weather" - )); - } - - #[tokio::test] - async fn tool_stream_generates_tool_call_id_when_parser_omits_one() { - let events = stream::iter(vec![Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - })]); - let parser = IdScriptedParser { - output: ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: "{}".to_string(), - }], - }, - tool_call_id: None, - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { id, name } - if id.starts_with("call_") && name == "get_weather" - )); - } - - #[tokio::test] - async fn real_buffered_parser_error_matches_streaming_and_non_streaming() { - let prefix = "I will check both.\n"; - let first_tool_call = concat!( - "<｜DSML｜tool_calls>\n", - "<｜DSML｜invoke name=\"get_weather\">\n", - "<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo\n", - "", - ); - let malformed_second_tool_call = concat!( - "\n<｜DSML｜invoke name=\"add\">\n", - "not a parameter\n", - "\n", - "", - ); - let streaming_chunks = vec![ - prefix.to_string(), - first_tool_call.to_string(), - malformed_second_tool_call.to_string(), - ]; - let full_output = streaming_chunks.concat(); - - let streaming = collect_deepseek_v4_message(streaming_chunks).await; - let non_streaming = collect_deepseek_v4_message(vec![full_output]).await; - - let expected = ( - format!("{prefix}{malformed_second_tool_call}"), - vec![( - "get_weather".to_string(), - serde_json::json!({ "location": "Tokyo" }), - )], - ); - assert_eq!(message_tool_projection(&streaming), expected); - assert_eq!(message_tool_projection(&non_streaming), expected); - } - - #[tokio::test] - async fn tool_parser_failure_falls_back_to_plain_text() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }), - Ok(ContentEvent::Done { - prompt_token_count: 3, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let collected = tool_event_stream( - events, - Some(Box::new(FailingParser { - fail_next: true, - buffered: String::new(), - })), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("tool stream should not fail"); - - assert_eq!( - events, - vec![ - AssistantEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }, - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }, - AssistantEvent::Done { - prompt_token_count: 3, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - - let message = ChatEventStream::new( - "req_fallback".to_string(), - Box::pin(structured_chat_event_stream(stream::iter( - events.into_iter().map(Ok), - ))), - ) - .collect_message() - .await - .expect("collect_message should succeed"); - assert_eq!(message.message.text(), "abcdef"); - assert!(message.message.tool_calls().next().is_none()); - } - - #[tokio::test] - async fn tool_stream_preserves_logprobs_delta() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.2, - rank: 1, - }], - }], - }), - token_ids: vec![], - }), - Ok(ContentEvent::Done { - prompt_token_count: 1, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - let events = tool_event_stream( - events, - Some(Box::new(FailingParser { - fail_next: false, - buffered: String::new(), - })), - ) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert_eq!( - events, - vec![ - AssistantEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - AssistantEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.2, - rank: 1, - }], - }], - }), - token_ids: vec![], - }, - AssistantEvent::Done { - prompt_token_count: 1, - output_token_count: 0, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - } - - #[tokio::test] - async fn tool_stream_rejects_interleaved_tool_indices() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - prompt_token_count: 1, - output_token_count: 1, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ToolParserOutput { - normal_text: String::new(), - calls: vec![ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: String::new(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 1, - name: None, - arguments: "{}".to_string(), - }, - ], - }], - finish_output: ToolParserOutput::default(), - }; - - let err = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .find_map(|output| output.err()) - .expect("expected invariant error"); - - assert!(matches!(err, Error::ToolCallStreamInvariant { .. })); - } - - #[tokio::test] - async fn tool_stream_resets_open_tool_when_normal_text_interrupts_it() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "start".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "text".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "args".to_string(), - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ - ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: None, - arguments: "}".to_string(), - }], - }, - ToolParserOutput { - normal_text: "plain text".to_string(), - calls: Vec::new(), - }, - ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: "{".to_string(), - }], - }, - ], - finish_output: ToolParserOutput::default(), - }; - - let err = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .find_map(|output| output.err()) - .expect("expected invariant error"); - - assert!(matches!( - err, - Error::ToolCallStreamInvariant { message } - if message == "received arguments for tool index 0 before any tool-call start" - )); - } - - #[tokio::test] - async fn tool_stream_emits_start_and_args_for_terminal_text() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - prompt_token_count: 1, - output_token_count: 1, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ToolParserOutput { - normal_text: String::new(), - calls: vec![ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: r#"{"a":1}"#.to_string(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 1, - name: Some("second".to_string()), - arguments: r#"{"b":2}"#.to_string(), - }, - ], - }], - finish_output: ToolParserOutput::default(), - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!(events[1], AssistantEvent::ToolCallStart { .. })); - assert!(matches!( - events[2], - AssistantEvent::ToolCallArgumentsDelta { .. } - )); - assert!(matches!(events[3], AssistantEvent::ToolCallStart { .. })); - assert!(matches!( - events[4], - AssistantEvent::ToolCallArgumentsDelta { .. } - )); - let collected = ChatEventStream::new( - "req_final_only".to_string(), - Box::pin(structured_chat_event_stream(stream::iter( - events.into_iter().map(Ok), - ))), - ) - .collect_message() - .await - .unwrap(); - let tool_calls = collected.message.tool_calls().collect::>(); - assert_eq!(tool_calls.len(), 2); - assert_eq!(tool_calls[0].name, "first"); - assert_eq!(tool_calls[1].name, "second"); - } -} diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs new file mode 100644 index 00000000000..038b907a26a --- /dev/null +++ b/rust/src/chat/src/output/default/unified.rs @@ -0,0 +1,685 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Adapts decoded text updates into parsed assistant deltas. +//! +//! This stage sits between low-level token decoding and final block assembly. +//! It drives one unified parser that may emit normal text, reasoning text, or +//! tool-call deltas, then normalizes those parser events into internal +//! assistant events. + +use asynk_strim_attr::{TryYielder, try_stream}; +use futures::{StreamExt as _, pin_mut}; +use thiserror_ext::AsReport; +use tracing::warn; +use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; +use vllm_text::output::DecodedTextEvent; + +use crate::Result; +use crate::error::Error; +use crate::event::AssistantBlockKind; +use crate::output::{AssistantEvent, DecodedTextEventStream, generate_tool_call_id}; + +/// Per-stream unified parsing state. +struct UnifiedParserState { + /// Parser for the current request stream. + parser: Box, + /// Whether unified parsing has already failed for this stream. + parser_failed: bool, + /// The parser-local index of the currently open tool call, if any. + /// + /// Supported parsers currently emit at most one active tool call at a time. + /// Change this to an indexed map if a model needs interleaved calls later. + open_call_index: Option, +} + +impl UnifiedParserState { + /// Create one fresh unified parsing state for a new streamed response. + fn new(parser: Box) -> Self { + Self { + parser, + parser_failed: false, + open_call_index: None, + } + } + + /// Initialize parser state once prompt token IDs are available. + fn initialize(&mut self, prompt_token_ids: &[u32]) { + if self.parser_failed { + return; + } + + match self.parser.initialize(prompt_token_ids) { + Ok(()) => {} + Err(error) => { + warn!( + error = %error.as_report(), + "failed to initialize unified parser; falling back to plain text deltas" + ); + self.parser_failed = true; + self.open_call_index = None; + } + } + } + + /// Convert one decoded text delta into zero or more parsed assistant events. + fn process_delta(&mut self, delta: String) -> Result> { + if self.parser_failed { + self.open_call_index = None; + return Ok(text_event(AssistantBlockKind::Text, delta).into_iter().collect()); + } + + let mut output = UnifiedParserOutput::default(); + match self.parser.parse_into(&delta, &mut output) { + Ok(()) => { + let mut events = Vec::new(); + self.process_parser_output(output, &mut events)?; + Ok(events) + } + Err(error) => { + warn!( + error = %error.as_report(), + "unified parser failed; falling back to plain text deltas" + ); + self.parser_failed = true; + + let mut events = Vec::new(); + self.process_parser_output(output, &mut events)?; + self.open_call_index = None; + + let recovered = self.parser.reset(); + if recovered.is_empty() && events.is_empty() { + push_text_delta(&mut events, AssistantBlockKind::Text, delta); + } else { + push_text_delta(&mut events, AssistantBlockKind::Text, recovered); + } + Ok(events) + } + } + } + + /// Flush parser state at end-of-stream and close any remaining open calls. + fn finish(&mut self) -> Result> { + let mut events = Vec::new(); + + if self.parser_failed { + return Ok(events); + } + + match self.parser.finish() { + Ok(output) => self.process_parser_output(output, &mut events)?, + Err(error) => { + warn!( + error = %error.as_report(), + "unified parser finish failed; closing open parser state" + ); + self.parser_failed = true; + self.open_call_index = None; + let recovered = self.parser.reset(); + push_text_delta(&mut events, AssistantBlockKind::Text, recovered); + } + } + + Ok(events) + } + + /// Apply one parsed unified output to the current stream state. + fn process_parser_output( + &mut self, + output: UnifiedParserOutput, + events: &mut Vec, + ) -> Result<()> { + for event in output.events { + match event { + UnifiedParserEvent::Text(delta) => { + self.open_call_index = None; + push_text_delta(events, AssistantBlockKind::Text, delta); + } + UnifiedParserEvent::Reasoning(delta) => { + self.open_call_index = None; + push_text_delta(events, AssistantBlockKind::Reasoning, delta); + } + UnifiedParserEvent::ToolCall(item) => { + self.process_tool_item(item, events)?; + } + } + } + + Ok(()) + } + + /// Apply one parsed tool-call delta emitted by the parser. + fn process_tool_item( + &mut self, + item: vllm_parser::tool::ToolCallDelta, + events: &mut Vec, + ) -> Result<()> { + if let Some(name) = item.name { + let is_new_tool = match self.open_call_index { + Some(open_call_index) => open_call_index != item.tool_index, + None => true, + }; + if is_new_tool { + let id = self + .parser + .tool_call_id(item.tool_index) + .map(str::to_string) + .unwrap_or_else(generate_tool_call_id); + self.open_call_index = Some(item.tool_index); + events.push(AssistantEvent::ToolCallStart { id, name }); + } + } + + if item.arguments.is_empty() { + return Ok(()); + } + let Some(open_call_index) = self.open_call_index else { + return Err(Error::ToolCallStreamInvariant { + message: format!( + "received arguments for tool index {} before any tool-call start", + item.tool_index + ), + }); + }; + if open_call_index != item.tool_index { + return Err(Error::ToolCallStreamInvariant { + message: format!( + "received arguments for tool index {} while tool index {} is open", + item.tool_index, open_call_index + ), + }); + } + + events.push(AssistantEvent::ToolCallArgumentsDelta { + delta: item.arguments, + }); + Ok(()) + } +} + +/// Build one plain text event if `delta` is non-empty. +fn text_event(kind: AssistantBlockKind, delta: String) -> Option { + if delta.is_empty() { + return None; + } + Some(AssistantEvent::TextDelta { kind, delta }) +} + +/// Push one plain text delta if it is non-empty. +fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { + if let Some(event) = text_event(kind, delta) { + events.push(event); + } +} + +/// Wrap one decoded-text stream into the internal unified assistant stream. +#[try_stream] +pub(crate) async fn unified_event_stream( + decoded_stream: impl DecodedTextEventStream, + parser: Box, + mut y: TryYielder, +) -> Result<()> { + pin_mut!(decoded_stream); + + let mut state = UnifiedParserState::new(parser); + + while let Some(event) = decoded_stream.next().await.transpose()? { + match event { + DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + } => { + state.initialize(&prompt_token_ids); + y.yield_ok(AssistantEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) + .await; + } + DecodedTextEvent::TextDelta { + delta, + token_ids, + logprobs, + finished, + } => { + for next in state.process_delta(delta)? { + y.yield_ok(next).await; + } + if logprobs.is_some() || !token_ids.is_empty() { + y.yield_ok(AssistantEvent::LogprobsDelta { + logprobs, + token_ids, + }) + .await; + } + if let Some(finished) = finished { + for next in state.finish()? { + y.yield_ok(next).await; + } + y.yield_ok(AssistantEvent::Done { + usage: finished.usage, + finish_reason: finished.finish_reason, + kv_transfer_params: finished.kv_transfer_params, + ec_transfer_params: finished.ec_transfer_params, + }) + .await; + } + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + use std::sync::Arc; + + use futures::{StreamExt as _, stream}; + use vllm_parser::reasoning::ReasoningError; + use vllm_parser::tool::{Tool, ToolCallDelta}; + use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput}; + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::unified_event_stream; + use crate::event::AssistantBlockKind; + use crate::output::AssistantEvent; + + enum ScriptedStep { + Output(UnifiedParserOutput), + Error { + committed: UnifiedParserOutput, + reset_text: String, + }, + } + + struct ScriptedParser { + steps: VecDeque, + reset_text: String, + tool_call_id: Option, + finish_error_reset_text: Option, + } + + impl ScriptedParser { + fn new(steps: impl IntoIterator) -> Self { + Self { + steps: steps.into_iter().collect(), + reset_text: String::new(), + tool_call_id: Some("call_test".to_string()), + finish_error_reset_text: None, + } + } + + fn with_finish_error(mut self, reset_text: &str) -> Self { + self.finish_error_reset_text = Some(reset_text.to_string()); + self + } + } + + impl vllm_parser::unified::UnifiedParser for ScriptedParser { + fn create( + _tools: &[vllm_parser::tool::Tool], + _tokenizer: vllm_tokenizer::DynTokenizer, + ) -> vllm_parser::unified::Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self::new([]))) + } + + fn parse_into( + &mut self, + _delta: &str, + output: &mut UnifiedParserOutput, + ) -> vllm_parser::unified::Result<()> { + match self.steps.pop_front().expect("unexpected parser call") { + ScriptedStep::Output(next) => { + output.append(next); + Ok(()) + } + ScriptedStep::Error { + committed, + reset_text, + } => { + output.append(committed); + self.reset_text = reset_text; + Err(UnifiedParserError::Reasoning( + ReasoningError::MissingToken { + token: "".to_string(), + }, + )) + } + } + } + + fn tool_call_id(&self, _tool_index: usize) -> Option<&str> { + self.tool_call_id.as_deref() + } + + fn finish(&mut self) -> vllm_parser::unified::Result { + if let Some(reset_text) = self.finish_error_reset_text.take() { + self.reset_text = reset_text; + return Err(UnifiedParserError::Reasoning( + ReasoningError::MissingToken { + token: "".to_string(), + }, + )); + } + Ok(UnifiedParserOutput::default()) + } + + fn reset(&mut self) -> String { + std::mem::take(&mut self.reset_text) + } + } + + fn decoded_delta(delta: &str) -> vllm_text::output::DecodedTextEvent { + vllm_text::output::DecodedTextEvent::TextDelta { + delta: delta.to_string(), + token_ids: Vec::new(), + logprobs: None, + finished: None, + } + } + + fn finished_delta(delta: &str) -> vllm_text::output::DecodedTextEvent { + vllm_text::output::DecodedTextEvent::TextDelta { + delta: delta.to_string(), + token_ids: Vec::new(), + logprobs: None, + finished: Some(vllm_text::output::Finished { + usage: vllm_llm::TokenUsage::default(), + finish_reason: crate::FinishReason::Stop(None), + kv_transfer_params: None, + ec_transfer_params: None, + }), + } + } + + async fn collect( + parser: ScriptedParser, + events: Vec, + ) -> Vec { + let stream = stream::iter(events.into_iter().map(Ok)); + unified_event_stream(stream, Box::new(parser)) + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap() + } + + fn text(delta: &str) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + output.push_text(delta.to_string()); + output + } + + fn reasoning(delta: &str) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + output.push_reasoning(delta.to_string()); + output + } + + fn tool_call(name: &str, arguments: &str) -> UnifiedParserOutput { + UnifiedParserOutput { + events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some(name.to_string()), + arguments: arguments.to_string(), + }, + )], + } + } + + fn tool_call_arguments(arguments: &str) -> UnifiedParserOutput { + UnifiedParserOutput { + events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall( + ToolCallDelta { + tool_index: 0, + name: None, + arguments: arguments.to_string(), + }, + )], + } + } + + fn combined(first: UnifiedParserOutput, second: UnifiedParserOutput) -> UnifiedParserOutput { + let mut output = first; + output.append(second); + output + } + + #[tokio::test] + async fn unified_stream_emits_reasoning_only_deltas() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(reasoning("thinking"))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![AssistantEvent::TextDelta { + kind: AssistantBlockKind::Reasoning, + delta: "thinking".to_string(), + }] + ); + } + + #[tokio::test] + async fn unified_stream_emits_tool_only_deltas() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(tool_call( + "get_weather", + r#"{"location":"Paris"}"#, + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_reasoning_followed_by_tool_call() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(combined( + reasoning("thinking"), + tool_call("get_weather", r#"{"location":"Paris"}"#), + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Reasoning, + delta: "thinking".to_string(), + }, + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_visible_text_followed_by_tool_call() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(combined( + text("visible "), + tool_call("get_weather", r#"{"location":"Paris"}"#), + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "visible ".to_string(), + }, + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_tool_arguments_before_trailing_text() { + let events = collect( + ScriptedParser::new([ + ScriptedStep::Output(tool_call("get_weather", "")), + ScriptedStep::Output(combined( + tool_call_arguments(r#"{"location":"Paris"}"#), + text(" done"), + )), + ]), + vec![decoded_delta("start"), decoded_delta("finish")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: " done".to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_fallback_keeps_committed_output_and_disables_later_parsing() { + let events = collect( + ScriptedParser::new([ScriptedStep::Error { + committed: text("committed"), + reset_text: "buffered".to_string(), + }]), + vec![decoded_delta("bad"), decoded_delta("later")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "committed".to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "buffered".to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "later".to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_finish_error_recovers_buffered_text() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())]) + .with_finish_error("buffered"), + vec![finished_delta("")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "buffered".to_string(), + }, + AssistantEvent::Done { + usage: vllm_llm::TokenUsage::default(), + finish_reason: crate::FinishReason::Stop(None), + kv_transfer_params: None, + ec_transfer_params: None, + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_recovers_incomplete_gemma4_tool_call_at_eos() { + let tokenizer = TestTokenizer::new() + .with_special_token("<|channel>", 256) + .with_special_token("", 257); + let tools = vec![Tool { + name: "write_file".to_string(), + description: None, + parameters: serde_json::json!({ "type": "object" }), + strict: None, + }]; + let parser = Gemma4UnifiedParser::new(&tools, Arc::new(tokenizer)).unwrap(); + let events = vec![ + decoded_delta("<|tool_call>"), + decoded_delta("call:write_file{"), + decoded_delta("content:<|\"|>hello "), + finished_delta("world<|\"|>"), + ]; + let stream = stream::iter(events.into_iter().map(Ok)); + let events = unified_event_stream(stream, Box::new(parser)) + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>" + .to_string(), + }, + AssistantEvent::Done { + usage: vllm_llm::TokenUsage::default(), + finish_reason: crate::FinishReason::Stop(None), + kv_transfer_params: None, + ec_transfer_params: None, + }, + ] + ); + } +} diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 5dc6bc31185..a84687882b1 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -1,19 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Native Harmony output processing for `gpt_oss`. //! //! Unlike the default text-first pipeline, this processor consumes //! `DecodedTextEvent` token IDs directly and lets the official `openai-harmony` //! parser recover the structured assistant message shape at token granularity. -use std::sync::LazyLock; - -use anyhow::Context; use asynk_strim_attr::{TryYielder, try_stream}; use futures::StreamExt as _; use openai_harmony::chat::{Content as HarmonyContent, Message as HarmonyMessage, Role}; -use openai_harmony::{ - HarmonyEncoding, HarmonyEncodingName, StreamableParser, load_harmony_encoding, -}; -use thiserror_ext::AsReport; +use openai_harmony::{HarmonyEncoding, StreamableParser}; use vllm_text::output::DecodedTextEvent; use crate::Result as ChatResult; @@ -24,6 +21,7 @@ use crate::output::{ generate_tool_call_id, }; use crate::parser::ParserSelection; +use crate::renderer::harmony::encoding::harmony_encoding; use crate::request::ChatRequest; /// Request-scoped Harmony output processor used for `model_type == "gpt_oss"`. @@ -35,6 +33,7 @@ use crate::request::ChatRequest; pub struct HarmonyChatOutputProcessor { encoding: &'static HarmonyEncoding, tool_calls_enabled: bool, + parallel_tool_calls: bool, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -76,6 +75,7 @@ impl HarmonyChatOutputProcessor { Ok(Self { encoding: harmony_encoding()?, tool_calls_enabled: request.tool_parsing_enabled(), + parallel_tool_calls: request.parallel_tool_calls, }) } } @@ -110,7 +110,11 @@ impl ChatOutputProcessor for HarmonyChatOutputProcessor { fn process(self: Box, decoded: DynDecodedTextEventStream) -> Result { let assistant = harmony_assistant_event_stream(decoded, self.encoding, self.tool_calls_enabled); - Ok(crate::output::structured::structured_chat_event_stream(assistant).boxed()) + Ok(crate::output::structured::structured_chat_event_stream( + assistant, + self.parallel_tool_calls, + ) + .boxed()) } } @@ -366,10 +370,10 @@ async fn harmony_assistant_event_stream( if let Some(finished) = finished { y.yield_ok(AssistantEvent::Done { - prompt_token_count: finished.prompt_token_count, - output_token_count: finished.output_token_count, + usage: finished.usage, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, + ec_transfer_params: finished.ec_transfer_params, }) .await; } @@ -379,18 +383,6 @@ async fn harmony_assistant_event_stream( Ok(()) } -/// Lazily load the shared GPT-OSS Harmony encoding once per process. -fn harmony_encoding() -> Result<&'static HarmonyEncoding> { - static ENCODING: LazyLock> = LazyLock::new(|| { - load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss) - .context("failed to load harmony encoding for gpt-oss") - }); - - ENCODING.as_ref().map_err(|error| Error::HarmonyOutputParsing { - error: error.to_report_string().into(), - }) -} - fn harmony_output_parsing_error( error: impl Into>, ) -> Error { diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index fe42542b473..9b9544bdc6b 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -1,14 +1,11 @@ -//! Harmony output tests share the upstream `openai-harmony` tiktoken cache. -//! -//! Use a file lock for tests that load the encoding so `cargo nextest` cannot -//! start multiple processes that concurrently populate the same cache file. +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::sync::Arc; use futures::executor::block_on; use futures::{TryStreamExt as _, stream}; use openai_harmony::chat::{Message, Role}; -use serial_test::file_serial; use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished}; use super::*; @@ -51,10 +48,14 @@ fn decoded_start() -> DecodedTextEvent { fn finished() -> Finished { Finished { - prompt_token_count: 0, - output_token_count: 0, + usage: vllm_llm::TokenUsage { + prompt_token_count: 0, + output_token_count: 0, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, } } @@ -88,7 +89,6 @@ fn request_with_tools() -> ChatRequest { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn interrupted_final_message_is_preserved() { let tokens = completion_tokens(&[text_message("final", "hello")]); let events = block_on(collect_events( @@ -112,16 +112,19 @@ fn interrupted_final_message_is_preserved() { text: "hello".to_string(), }], }, - prompt_token_count: 0, - output_token_count: 0, + usage: vllm_llm::TokenUsage { + prompt_token_count: 0, + output_token_count: 0, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }) ); } #[test] -#[file_serial(harmony_tiktoken_cache)] fn eos_flush_preserves_trailing_replacement_text() { let mut tokens = completion_tokens(&[text_message("final", "Hi")]); tokens.pop(); @@ -147,7 +150,6 @@ fn eos_flush_preserves_trailing_replacement_text() { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn interrupted_analysis_message_is_preserved() { let tokens = completion_tokens(&[text_message("analysis", "think")]); let events = block_on(collect_events( @@ -171,16 +173,19 @@ fn interrupted_analysis_message_is_preserved() { text: "think".to_string(), }], }, - prompt_token_count: 0, - output_token_count: 0, + usage: vllm_llm::TokenUsage { + prompt_token_count: 0, + output_token_count: 0, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }) ); } #[test] -#[file_serial(harmony_tiktoken_cache)] fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() { let tokens = completion_tokens(&[ text_message("commentary", "Let me check."), @@ -208,7 +213,6 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn multiple_messages_get_newline_separators() { let tokens = completion_tokens(&[ text_message("analysis", "first think"), @@ -240,7 +244,6 @@ fn multiple_messages_get_newline_separators() { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn tool_calls_stream_arguments_and_finish_with_local_id_shape() { let tokens = completion_tokens(&[tool_message( "get_weather", @@ -293,7 +296,6 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn semantic_events_precede_same_update_logprobs() { let tokens = completion_tokens(&[text_message("final", "hello")]); let events = block_on(collect_events( @@ -344,7 +346,6 @@ fn rejects_generic_parser_overrides() { } #[test] -#[file_serial(harmony_tiktoken_cache)] fn allows_auto_auto_only() { validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap(); let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(); diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 6dda8ba0dae..ecf90015caa 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -1,10 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::pin::Pin; use std::sync::Arc; use futures::Stream; -use subenum::subenum; use trait_set::trait_set; use uuid::Uuid; +use vllm_llm::TokenUsage; use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent}; use crate::FinishReason; @@ -21,23 +24,19 @@ pub(crate) use harmony::validate_harmony_parser_overrides; /// Internal assistant event before final assembly. /// -/// - [`ContentEvent`]: subenum after reasoning parsing, carries only text content. -/// - [`AssistantEvent`]: full event after tool parsing, adds tool-call variants. -#[subenum(ContentEvent)] +/// Unified parsing produces these events, and structured assembly consumes +/// them to build public chat events. #[derive(Debug, Clone, PartialEq)] pub(crate) enum AssistantEvent { - #[subenum(ContentEvent)] Start { prompt_token_ids: Arc<[u32]>, prompt_logprobs: Option, }, - #[subenum(ContentEvent)] TextDelta { kind: AssistantBlockKind, delta: String, }, /// Per-decoded-update sample metadata: logprobs and/or output token IDs. - #[subenum(ContentEvent)] LogprobsDelta { logprobs: Option, token_ids: Vec, @@ -47,61 +46,17 @@ pub(crate) enum AssistantEvent { /// A delta for the arguments of the currently open tool call. Must follow a /// `ToolCallStart`. ToolCallArgumentsDelta { delta: String }, - #[subenum(ContentEvent)] Done { - prompt_token_count: usize, - output_token_count: usize, + usage: TokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. kv_transfer_params: Option, + /// Connector-specific encoder cache transfer parameters for + /// disaggregated serving. + ec_transfer_params: Option, }, } -impl ContentEvent { - /// Convert a [`DecodedTextEvent`] into one or more [`ContentEvent`] values - /// by treating all text as plain (non-reasoning) content. - fn from_decoded_plain_text(event: DecodedTextEvent) -> Vec { - match event { - DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => vec![Self::Start { - prompt_token_ids, - prompt_logprobs, - }], - DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - } => { - let mut events = Vec::new(); - if !delta.is_empty() { - events.push(Self::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }); - } - if logprobs.is_some() || !token_ids.is_empty() { - events.push(Self::LogprobsDelta { - logprobs, - token_ids, - }); - } - if let Some(finished) = finished { - events.push(Self::Done { - prompt_token_count: finished.prompt_token_count, - output_token_count: finished.output_token_count, - finish_reason: finished.finish_reason, - kv_transfer_params: finished.kv_transfer_params, - }); - } - events - } - } - } -} - /// Boxed stream of decoded text events coming from [`vllm_text`]. pub type DynDecodedTextEventStream = Pin> + Send>>; /// Boxed stream of structured chat events exposed by [`crate::ChatLlm`]. diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index ed6e3a5130c..387d8f9bd53 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Adapts parsed assistant updates into structured chat events. //! //! This module remains the final assembly stage in `vllm-chat`. Token-to-text @@ -53,16 +56,22 @@ struct StructuredEventState { open_tool_call: Option, /// Next OpenAI-compatible tool-call ordinal. next_tool_call_index: usize, + /// Whether more than one tool call may be surfaced northbound. + parallel_tool_calls: bool, + /// Whether the current tool-call parse is being suppressed. + suppressing_tool_call: bool, } impl StructuredEventState { /// Create one fresh assembly state for a new streamed response. - fn new() -> Self { + fn new(parallel_tool_calls: bool) -> Self { Self { message: AssistantMessage::default(), open_text_block: None, open_tool_call: None, next_tool_call_index: 0, + parallel_tool_calls, + suppressing_tool_call: false, } } @@ -98,6 +107,12 @@ impl StructuredEventState { let index = self.next_tool_call_index; self.next_tool_call_index += 1; + if !self.parallel_tool_calls && index >= 1 { + self.suppressing_tool_call = true; + return Ok(events); + } + + self.suppressing_tool_call = false; self.open_tool_call = Some(OpenToolCall { index, id: id.clone(), @@ -110,6 +125,10 @@ impl StructuredEventState { /// Append one incremental tool-call arguments delta. fn push_tool_call_arguments(&mut self, delta: String) -> Result> { + if self.suppressing_tool_call { + return Ok(Vec::new()); + } + let mut events = Vec::new(); let Some(open_tool_call) = self.open_tool_call.as_mut() else { return Err(Error::ToolCallStreamInvariant { @@ -127,20 +146,20 @@ impl StructuredEventState { /// Close any open block and emit the terminal `Done` event. fn finish( &mut self, - prompt_token_count: usize, - output_token_count: usize, + usage: vllm_llm::TokenUsage, finish_reason: FinishReason, kv_transfer_params: Option, + ec_transfer_params: Option, ) -> Result> { let mut events = Vec::new(); self.close_open_text_block(&mut events); self.close_open_tool_call(&mut events); events.push(ChatEvent::Done { message: self.message.clone(), - prompt_token_count, - output_token_count, + usage, finish_reason, kv_transfer_params, + ec_transfer_params, }); Ok(events) } @@ -209,6 +228,11 @@ impl StructuredEventState { /// Finalize the currently open tool call, if present. fn close_open_tool_call(&mut self, events: &mut Vec) { + if self.suppressing_tool_call { + self.suppressing_tool_call = false; + return; + } + let Some(open_tool_call) = self.open_tool_call.take() else { return; }; @@ -231,11 +255,12 @@ impl StructuredEventState { #[try_stream] pub(crate) async fn structured_chat_event_stream( stream: impl AssistantEventStream, + parallel_tool_calls: bool, mut y: TryYielder, ) -> Result<()> { pin_mut!(stream); - let mut state = StructuredEventState::new(); + let mut state = StructuredEventState::new(parallel_tool_calls); while let Some(event) = stream.next().await.transpose()? { match event { @@ -273,17 +298,14 @@ pub(crate) async fn structured_chat_event_stream( } } AssistantEvent::Done { - prompt_token_count, - output_token_count, + usage, finish_reason, kv_transfer_params, + ec_transfer_params, } => { - for next in state.finish( - prompt_token_count, - output_token_count, - finish_reason, - kv_transfer_params, - )? { + for next in + state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)? + { y.yield_ok(next).await; } } @@ -313,14 +335,18 @@ mod tests { delta: r#"{"city":"Paris"}"#.to_string(), }), Ok(AssistantEvent::Done { - prompt_token_count: 1, - output_token_count: 1, + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); - let events = structured_chat_event_stream(events) + let events = structured_chat_event_stream(events, true) .collect::>() .await .into_iter() @@ -364,14 +390,18 @@ mod tests { delta: r#"{"b":2}"#.to_string(), }), Ok(AssistantEvent::Done { - prompt_token_count: 1, - output_token_count: 1, + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); - let events = structured_chat_event_stream(events) + let events = structured_chat_event_stream(events, true) .collect::>() .await .into_iter() @@ -412,14 +442,18 @@ mod tests { delta: r#"{"city":"Paris"}"#.to_string(), }), Ok(AssistantEvent::Done { - prompt_token_count: 1, - output_token_count: 1, + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); - let events = structured_chat_event_stream(events) + let events = structured_chat_event_stream(events, true) .collect::>() .await .into_iter() @@ -460,14 +494,18 @@ mod tests { delta: "done".to_string(), }), Ok(AssistantEvent::Done { - prompt_token_count: 1, - output_token_count: 1, + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, + ec_transfer_params: None, }), ]); - let events = structured_chat_event_stream(events) + let events = structured_chat_event_stream(events, true) .collect::>() .await .into_iter() @@ -495,7 +533,7 @@ mod tests { delta: "{}".to_string(), })]); - let err = structured_chat_event_stream(events) + let err = structured_chat_event_stream(events, true) .collect::>() .await .into_iter() @@ -505,4 +543,57 @@ mod tests { assert!(matches!(err, Error::ToolCallStreamInvariant { .. })); } + + #[tokio::test] + async fn structured_stream_suppresses_later_tool_calls_when_parallel_disabled() { + let events = stream::iter(vec![ + Ok(AssistantEvent::ToolCallStart { + id: "call_1".to_string(), + name: "first".to_string(), + }), + Ok(AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"a":1}"#.to_string(), + }), + Ok(AssistantEvent::ToolCallStart { + id: "call_2".to_string(), + name: "second".to_string(), + }), + Ok(AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"b":2}"#.to_string(), + }), + Ok(AssistantEvent::Done { + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, + finish_reason: FinishReason::stop_eos(), + kv_transfer_params: None, + ec_transfer_params: None, + }), + ]); + + let events = structured_chat_event_stream(events, false) + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap(); + + assert!(matches!( + events[0], + ChatEvent::ToolCallStart { index: 0, .. } + )); + assert!(matches!( + events[1], + ChatEvent::ToolCallArgumentsDelta { index: 0, .. } + )); + assert!(matches!(events[2], ChatEvent::ToolCallEnd { index: 0, .. })); + let ChatEvent::Done { message, .. } = &events[3] else { + panic!("expected done"); + }; + let tool_calls = message.tool_calls().collect::>(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].name, "first"); + } } diff --git a/rust/src/chat/src/parser/mod.rs b/rust/src/chat/src/parser/mod.rs index 244a87cc7a7..7684aa9d973 100644 --- a/rust/src/chat/src/parser/mod.rs +++ b/rust/src/chat/src/parser/mod.rs @@ -1,5 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + pub mod reasoning; pub mod tool; +pub mod unified; use std::collections::HashMap; use std::convert::Infallible; diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index 09111d7252f..10d48b2b395 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -1,12 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Reasoning parser registration and selection boundary for `vllm-chat`. -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; -pub use vllm_reasoning_parser::{ +pub use vllm_parser::reasoning::{ CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser, - DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser, - KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser, - ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser, + DeepSeekV4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, + MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, NemotronV3ReasoningParser, + Qwen3ReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, + Step3ReasoningParser, Step3p5ReasoningParser, }; use vllm_tokenizer::DynTokenizer; @@ -23,14 +27,18 @@ pub mod names { pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; pub const MINIMAX_M2: &str = "minimax_m2"; + pub const MINIMAX_M3: &str = "minimax_m3"; pub const NEMOTRON_V3: &str = "nemotron_v3"; pub const QWEN3: &str = "qwen3"; + pub const SEED_OSS: &str = "seed_oss"; pub const STEP3: &str = "step3"; + pub const STEP3P5: &str = "step3p5"; } /// Constructor signature for one registered reasoning parser implementation. -type ReasoningParserCreator = - fn(DynTokenizer) -> vllm_reasoning_parser::Result>; +type ReasoningParserCreator = Arc< + dyn Fn(DynTokenizer) -> vllm_parser::reasoning::Result> + Send + Sync, +>; /// Registry and model matcher for reasoning parsers. pub type ReasoningParserFactory = ParserFactory; @@ -54,14 +62,17 @@ impl ReasoningParserFactory { .register_parser::(names::DEEPSEEK_R1) .register_parser::(names::DEEPSEEK_V3) .register_parser::(names::DEEPSEEK_V4) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) .register_parser::(names::MINIMAX_M2) + .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) .register_parser::(names::QWEN3) - .register_parser::(names::STEP3); + .register_parser::(names::SEED_OSS) + .register_parser::(names::STEP3) + .register_parser::(names::STEP3P5); factory .register_pattern("deepseek-r1", names::DEEPSEEK_R1) @@ -77,7 +88,16 @@ impl ReasoningParserFactory { .register_pattern("glm-4.5", names::GLM45) .register_pattern("kimi-k2", names::KIMI_K2) .register_pattern("kimi", names::KIMI) + // step3p5 patterns must precede `step3`: substring matching would + // otherwise route step3p5 IDs to step3. + .register_pattern("step-3p5", names::STEP3P5) + .register_pattern("step3p5", names::STEP3P5) + .register_pattern("step-3.5", names::STEP3P5) .register_pattern("step3", names::STEP3) + .register_pattern("seed-oss", names::SEED_OSS) + .register_pattern("seedoss", names::SEED_OSS) + .register_pattern("minimax-m3", names::MINIMAX_M3) + .register_pattern("mm-m3", names::MINIMAX_M3) .register_pattern("minimax", names::MINIMAX_M2) .register_pattern("mm-m2", names::MINIMAX_M2) .register_pattern("cohere", names::COHERE_CMD) @@ -93,7 +113,17 @@ impl ReasoningParserFactory { where T: ReasoningParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split reasoning registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ReasoningError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -108,7 +138,7 @@ impl ReasoningParserFactory { available_names: self.list(), })?; - creator(tokenizer).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tokenizer).map_err(|error| crate::Error::ParserInitialization { kind: "reasoning", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/reasoning/tests.rs b/rust/src/chat/src/parser/reasoning/tests.rs index 89b5f8e2308..7b7f1fe94a5 100644 --- a/rust/src/chat/src/parser/reasoning/tests.rs +++ b/rust/src/chat/src/parser/reasoning/tests.rs @@ -1,39 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::sync::Arc; -use vllm_tokenizer::Tokenizer; +use vllm_tokenizer::test_utils::TestTokenizer; use super::{ReasoningParserFactory, names}; -struct FakeTokenizer; - -impl Tokenizer for FakeTokenizer { - fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result> { - Ok(text.chars().map(u32::from).collect()) - } - - fn decode( - &self, - token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_tokenizer::Result { - Ok(token_ids - .iter() - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) - .collect()) - } - - fn token_to_id(&self, _token: &str) -> Option { - None - } -} - #[test] fn factory_contains_and_lists_registered_parsers() { let factory = ReasoningParserFactory::new(); assert!(factory.contains(names::QWEN3)); assert!(factory.contains(names::DEEPSEEK_V4)); + assert!(factory.contains(names::SEED_OSS)); + assert!(factory.contains(names::STEP3P5)); + assert!(factory.contains(names::MINIMAX_M3)); + assert!(factory.contains(names::GEMMA4)); assert!(factory.list().contains(&names::QWEN3.to_string())); assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string())); + assert!(factory.list().contains(&names::SEED_OSS.to_string())); + assert!(factory.list().contains(&names::STEP3P5.to_string())); + assert!(factory.list().contains(&names::MINIMAX_M3.to_string())); + assert!(factory.list().contains(&names::GEMMA4.to_string())); } #[test] @@ -49,9 +37,57 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() { ); } +#[test] +fn factory_routes_step3p5_models_to_dedicated_parser() { + let factory = ReasoningParserFactory::new(); + // step3p5 patterns must beat the bare `step3` substring. + assert_eq!( + factory.resolve_name_for_model("step-3p5-instruct"), + Some(names::STEP3P5) + ); + assert_eq!( + factory.resolve_name_for_model("step3p5"), + Some(names::STEP3P5) + ); + assert_eq!( + factory.resolve_name_for_model("step-3.5-base"), + Some(names::STEP3P5) + ); + assert_eq!( + factory.resolve_name_for_model("step3-base"), + Some(names::STEP3) + ); +} + +#[test] +fn factory_routes_seed_oss_models() { + let factory = ReasoningParserFactory::new(); + assert_eq!( + factory.resolve_name_for_model("ByteDance-Seed/Seed-OSS-36B-Instruct"), + Some(names::SEED_OSS) + ); + assert_eq!( + factory.resolve_name_for_model("seedoss-7b"), + Some(names::SEED_OSS) + ); +} + +#[test] +fn factory_resolves_minimax_m3_before_generic_minimax() { + let factory = ReasoningParserFactory::new(); + assert_eq!( + factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"), + Some(names::MINIMAX_M3) + ); + assert_eq!( + factory.resolve_name_for_model("mm-m3"), + Some(names::MINIMAX_M3) + ); +} + #[test] fn factory_rejects_unknown_parser_names() { - let tokenizer = Arc::new(FakeTokenizer); + let tokenizer = Arc::new(TestTokenizer::new()); let factory = ReasoningParserFactory::new(); let error = match factory.create("missing", tokenizer) { Ok(_) => panic!("expected parser lookup to fail"), diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 29961d1d82a..3e394c417bc 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -1,13 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! Tool parser registration and selection boundary for `vllm-chat`. -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; -pub use vllm_tool_parser::{ +pub use vllm_parser::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, - Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser, + Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, - MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, - ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput, + MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, + Qwen3XmlToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -22,6 +25,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const GLM47: &str = "glm47"; pub const GEMMA4: &str = "gemma4"; + pub const GRANITE4: &str = "granite4"; pub const HERMES: &str = "hermes"; pub const HY_V3: &str = "hy_v3"; // Matches the Python CLI name `--tool-call-parser internlm`, which Python @@ -31,6 +35,7 @@ pub mod names { pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; pub const MINIMAX_M2: &str = "minimax_m2"; + pub const MINIMAX_M3: &str = "minimax_m3"; pub const MISTRAL: &str = "mistral"; pub const PHI4_MINI_JSON: &str = "phi4_mini_json"; pub const QWEN3_CODER: &str = "qwen3_coder"; @@ -38,7 +43,8 @@ pub mod names { } /// Constructor signature for one registered tool parser implementation. -type ToolParserCreator = fn(&[ChatTool]) -> vllm_tool_parser::Result>; +type ToolParserCreator = + Arc vllm_parser::tool::Result> + Send + Sync>; /// Registry and model matcher for tool parsers. pub type ToolParserFactory = ParserFactory; @@ -63,7 +69,8 @@ impl ToolParserFactory { .register_parser::(names::DEEPSEEK_V4) .register_parser::(names::GLM45) .register_parser::(names::GLM47) - .register_parser::(names::GEMMA4) + .register_unified_dummy(names::GEMMA4) + .register_parser::(names::GRANITE4) .register_parser::(names::HERMES) .register_parser::(names::HY_V3) .register_parser::(names::INTERNLM) @@ -71,6 +78,7 @@ impl ToolParserFactory { .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) + .register_parser::(names::MINIMAX_M3) .register_parser::(names::MISTRAL) .register_parser::(names::PHI4_MINI_JSON) .register_parser::(names::QWEN3_XML) @@ -107,7 +115,10 @@ impl ToolParserFactory { .register_pattern("glm-4.5", names::GLM45) .register_pattern("gemma4", names::GEMMA4) .register_pattern("gemma-4", names::GEMMA4) + .register_pattern("granite-4", names::GRANITE4) .register_pattern("kimi-k2", names::KIMI_K2) + .register_pattern("minimax-m3", names::MINIMAX_M3) + .register_pattern("mm-m3", names::MINIMAX_M3) .register_pattern("minimax", names::MINIMAX_M2) .register_pattern("mm-m2", names::MINIMAX_M2); @@ -119,7 +130,17 @@ impl ToolParserFactory { where T: ToolParser + 'static, { - self.register_creator(name, T::create) + self.register_creator(name, Arc::new(T::create)) + } + + /// Register one unified-only parser name in the split tool registry. + pub fn register_unified_dummy(&mut self, name: &str) -> &mut Self { + let name = name.to_string(); + let registered_name = name.clone(); + self.register_creator( + ®istered_name, + Arc::new(move |_| Err(ToolParserError::DummyUnifiedParser { name: name.clone() })), + ) } /// Construct a parser from an exact name. @@ -130,7 +151,7 @@ impl ToolParserFactory { available_names: self.list(), })?; - creator(tools).map_err(|error| crate::Error::ParserInitialization { + creator.as_ref()(tools).map_err(|error| crate::Error::ParserInitialization { kind: "tool", name: name.to_string(), error: error.into(), diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index 6fd380bd223..bb89ea38feb 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -1,6 +1,9 @@ -use vllm_tool_parser::Result; +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use super::{ToolParser, ToolParserFactory, ToolParserOutput, names}; +use vllm_parser::tool::{Result, ToolParserOutput}; + +use super::{ToolParser, ToolParserFactory, names}; use crate::Error; use crate::request::ChatTool; @@ -145,6 +148,10 @@ fn factory_new_resolves_default_patterns() { factory.resolve_name_for_model("google/gemma-4-27b-it"), Some(names::GEMMA4) ); + assert_eq!( + factory.resolve_name_for_model("ibm-granite/granite-4.0-h-tiny"), + Some(names::GRANITE4) + ); assert_eq!( factory.resolve_name_for_model("NousResearch/Hermes-3-Llama-3.1-8B"), Some(names::HERMES) @@ -153,6 +160,14 @@ fn factory_new_resolves_default_patterns() { factory.resolve_name_for_model("tencent/Hy3-preview"), Some(names::HY_V3) ); + assert_eq!( + factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"), + Some(names::MINIMAX_M3) + ); + assert_eq!( + factory.resolve_name_for_model("org/mm-m3-base"), + Some(names::MINIMAX_M3) + ); assert_eq!( factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"), Some(names::MINIMAX_M2) diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs new file mode 100644 index 00000000000..ff97bb38db1 --- /dev/null +++ b/rust/src/chat/src/parser/unified.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser registration and selection boundary for `vllm-chat`. + +use std::sync::LazyLock; + +pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser}; +use vllm_tokenizer::DynTokenizer; + +use crate::parser::ParserFactory; +use crate::request::ChatTool; + +/// Canonical public names for registered unified parsers. +pub mod names { + pub const GEMMA4: &str = "gemma4"; +} + +/// Constructor signature for one registered unified parser implementation. +type UnifiedParserCreator = + fn(&[ChatTool], DynTokenizer) -> vllm_parser::unified::Result>; + +/// Registry and model matcher for unified parsers. +pub type UnifiedParserFactory = ParserFactory; + +impl UnifiedParserFactory { + /// Get the global unified parser factory with built-in registrations and + /// model mappings. + pub fn global() -> &'static Self { + static INSTANCE: LazyLock = LazyLock::new(UnifiedParserFactory::new); + &INSTANCE + } + + /// Create the default registry with built-in parser names and model + /// mappings. + pub fn new() -> Self { + let mut factory = Self::default(); + + factory.register_parser::(names::GEMMA4); + + factory + .register_pattern("gemma-4", names::GEMMA4) + .register_pattern("gemma4", names::GEMMA4); + + factory + } + + /// Register one parser type that exposes a static `create()` constructor. + pub fn register_parser(&mut self, name: &str) -> &mut Self + where + T: UnifiedParser + 'static, + { + self.register_creator(name, T::create) + } + + /// Construct a parser from an exact name. + pub fn create( + &self, + name: &str, + tools: &[ChatTool], + tokenizer: DynTokenizer, + ) -> crate::Result> { + let creator = self.creator(name).ok_or_else(|| crate::Error::ParserUnavailableByName { + kind: "unified", + name: name.to_string(), + available_names: self.list(), + })?; + + creator(tools, tokenizer).map_err(|error| crate::Error::ParserInitialization { + kind: "unified", + name: name.to_string(), + error: error.into(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::{UnifiedParserFactory, names}; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|channel>", 256) + .with_regular_token("", 257) + } + + #[test] + fn factory_registers_gemma4() { + let factory = UnifiedParserFactory::new(); + + assert!(factory.contains(names::GEMMA4)); + assert_eq!( + factory.resolve_name_for_model("google/gemma-4-27b-it"), + Some(names::GEMMA4) + ); + factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap(); + } +} diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 97825519276..86925c87c4c 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! DeepSeek V3.2 prompt renderer. use std::collections::{HashMap, HashSet}; @@ -49,6 +52,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { let last_user_render_index = find_last_user_render_index(request.messages.as_slice(), render_offset); let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice()); + let continue_final_message = request.chat_options.continue_final_message(); let mut prompt = String::from(BOS_TOKEN); if request.tool_parsing_enabled() { @@ -66,6 +70,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { last_user_actual_index, thinking_mode, drop_thinking, + continue_final_message, )?; } @@ -96,6 +101,7 @@ fn render_message( last_user_actual_index: usize, thinking_mode: ThinkingMode, drop_thinking: bool, + continue_final_message: bool, ) -> Result<()> { let render_index = message_index as isize + render_offset; let opens_thinking = render_index == last_user_render_index; @@ -125,9 +131,7 @@ fn render_message( thinking_mode, drop_thinking, ), - // TODO: Respect `continue_final_message` and map it to DeepSeek's - // prefix-style final-assistant continuation behavior. - false, + continue_final_message && message_index + 1 == messages.len(), ), ChatMessage::ToolResponse { content, .. } => render_tool_message( out, diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs index 97225bbab09..1a41e14f8fc 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs @@ -1,8 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + mod encoding; use vllm_text::Prompt; -use super::{ChatRenderer, RenderedPrompt}; +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; use crate::Result; use crate::request::ChatRequest; @@ -23,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(encoding::render_request(request)?), + effective_template_kwargs: request_template_kwargs(request), }) } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 0b8f2b09e11..76796edd096 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -1,82 +1,21 @@ -use std::fs; +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::path::PathBuf; use expect_test::{ExpectFile, expect, expect_file}; -use serde::Deserialize; use serde_json::{Value, json}; use thiserror_ext::AsReport; use super::DeepSeekV32ChatRenderer; use crate::error::Error; use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, }; use crate::{ChatRenderer, ChatRole}; -#[derive(Debug, Deserialize)] -struct FixtureRequest { - #[serde(default)] - tools: Vec, - messages: Vec, -} - -#[derive(Debug, Deserialize)] -struct FixtureTool { - function: FixtureToolFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolFunction { - name: String, - description: Option, - parameters: Value, - #[serde(default)] - strict: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "role", rename_all = "snake_case")] -enum FixtureMessage { - System { - content: String, - }, - Developer { - content: String, - #[serde(default)] - tools: Vec, - }, - User { - content: String, - }, - Assistant { - #[serde(default)] - content: String, - #[serde(default)] - reasoning_content: String, - #[serde(default)] - tool_calls: Vec, - }, - Tool { - content: String, - #[serde(default)] - tool_call_id: Option, - }, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCall { - #[serde(default)] - id: Option, - function: FixtureToolCallFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCallFunction { - name: String, - arguments: String, -} - fn render_request(request: &ChatRequest) -> String { DeepSeekV32ChatRenderer::new() .render(request) @@ -115,88 +54,14 @@ fn thinking_request(messages: Vec) -> ChatRequest { } fn fixture_request(input_name: &str) -> ChatRequest { - let fixture = fs::read_to_string(fixture_path(input_name)).unwrap(); - let fixture: FixtureRequest = serde_json::from_str(&fixture).unwrap(); - let mut request = ChatRequest { - request_id: "deepseek-v32-fixture".to_string(), - messages: fixture - .messages - .into_iter() - .enumerate() - .map(|(index, message)| match message { - FixtureMessage::System { content } => ChatMessage::system(content), - FixtureMessage::Developer { content, tools } => ChatMessage::developer( - content, - (!tools.is_empty()).then(|| to_chat_tools(&tools)), - ), - FixtureMessage::User { content } => ChatMessage::user(content), - FixtureMessage::Assistant { - content, - reasoning_content, - tool_calls, - } => { - let mut blocks = Vec::new(); - if !reasoning_content.is_empty() { - blocks.push(AssistantContentBlock::Reasoning { - text: reasoning_content, - }); - } - if !content.is_empty() { - blocks.push(AssistantContentBlock::Text { text: content }); - } - blocks.extend(tool_calls.into_iter().enumerate().map( - |(tool_index, tool_call)| { - AssistantContentBlock::ToolCall(AssistantToolCall { - id: tool_call.id.unwrap_or_else(|| { - format!("fixture-tool-call-{index}-{tool_index}") - }), - name: tool_call.function.name, - arguments: tool_call.function.arguments, - }) - }, - )); - ChatMessage::assistant_blocks(blocks) - } - FixtureMessage::Tool { - content, - tool_call_id, - } => ChatMessage::tool_response( - content, - tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")), - ), - }) - .collect(), - tools: to_chat_tools(&fixture.tools), - tool_choice: if fixture.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, - ..ChatRequest::for_test() - }; - if matches!( - request.messages.last().map(ChatMessage::role), - Some(ChatRole::Assistant) - ) { - request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; - } - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), Value::Bool(true)); - request + fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options()) } -fn to_chat_tools(tools: &[FixtureTool]) -> Vec { - tools - .iter() - .map(|tool| ChatTool { - name: tool.function.name.clone(), - description: tool.function.description.clone(), - parameters: tool.function.parameters.clone(), - strict: tool.function.strict, - }) - .collect() +fn deepseek_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + enable_thinking: true, + no_generation_prompt_when_last_assistant: true, + } } fn fixture_path(name: &str) -> PathBuf { @@ -404,6 +269,24 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() { expect!["chat template error: invalid DeepSeek V3.2 assistant message after last user message: expected reasoning or tool calls"] .assert_eq(&error.to_report_string()); } + +#[test] +fn continue_final_assistant_omits_final_eos() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("write"), + ChatMessage::assistant_text("partial answer"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜>partial answer"] + .assert_eq(&rendered); +} + #[test] fn render_rejects_multimodal_input() { let request = ChatRequest { diff --git a/rust/src/chat/src/renderer/deepseek_v4/encoding.rs b/rust/src/chat/src/renderer/deepseek_v4/encoding.rs index 54a69248618..80438b0c415 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/encoding.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + //! DeepSeek V4 prompt renderer. //! //! Original Python implementation: diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs index 7c3f4631d20..2762f1546f3 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs @@ -1,8 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + mod encoding; use vllm_text::Prompt; -use super::{ChatRenderer, RenderedPrompt}; +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; use crate::Result; use crate::request::ChatRequest; @@ -22,6 +25,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(encoding::render_request(request)?), + effective_template_kwargs: request_template_kwargs(request), }) } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 78936d8e68e..73380cef260 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -1,95 +1,16 @@ -use std::fs; +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::path::PathBuf; use expect_test::{ExpectFile, expect, expect_file}; -use serde::Deserialize; use serde_json::Value; use super::DeepSeekV4ChatRenderer; +use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; -use crate::request::{ - ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, ReasoningEffort, -}; -use crate::{ChatRenderer, ChatRole}; - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -enum FixtureFile { - WithTools(FixtureRequest), - MessagesOnly(Vec), -} - -#[derive(Debug, Deserialize)] -struct FixtureRequest { - #[serde(default)] - tools: Vec, - messages: Vec, -} - -impl FixtureFile { - fn into_parts(self) -> (Vec, Vec) { - match self { - Self::WithTools(req) => (req.tools, req.messages), - Self::MessagesOnly(messages) => (Vec::new(), messages), - } - } -} - -#[derive(Debug, Deserialize)] -struct FixtureTool { - function: FixtureToolFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolFunction { - name: String, - description: Option, - parameters: Value, - #[serde(default)] - strict: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "role", rename_all = "snake_case")] -enum FixtureMessage { - System { - content: String, - }, - Developer { - content: String, - #[serde(default)] - tools: Vec, - }, - User { - content: String, - }, - Assistant { - #[serde(default)] - content: String, - #[serde(default)] - reasoning_content: String, - #[serde(default)] - tool_calls: Vec, - }, - Tool { - content: String, - #[serde(default)] - tool_call_id: Option, - }, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCall { - #[serde(default)] - id: Option, - function: FixtureToolCallFunction, -} - -#[derive(Debug, Deserialize)] -struct FixtureToolCallFunction { - name: String, - arguments: String, -} +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort}; fn render_request(request: &ChatRequest) -> String { DeepSeekV4ChatRenderer::new() @@ -101,88 +22,14 @@ fn render_request(request: &ChatRequest) -> String { } fn fixture_request(input_name: &str) -> ChatRequest { - let fixture = fs::read_to_string(fixture_path(input_name)).unwrap(); - let fixture: FixtureFile = serde_json::from_str(&fixture).unwrap(); - let (fixture_tools, fixture_messages) = fixture.into_parts(); - let mut request = ChatRequest { - request_id: "deepseek-v4-fixture".to_string(), - messages: fixture_messages - .into_iter() - .enumerate() - .map(|(index, message)| match message { - FixtureMessage::System { content } => ChatMessage::system(content), - FixtureMessage::Developer { content, tools } => ChatMessage::developer( - content, - (!tools.is_empty()).then(|| to_chat_tools(&tools)), - ), - FixtureMessage::User { content } => ChatMessage::user(content), - FixtureMessage::Assistant { - content, - reasoning_content, - tool_calls, - } => { - let mut blocks = Vec::new(); - if !reasoning_content.is_empty() { - blocks.push(AssistantContentBlock::Reasoning { - text: reasoning_content, - }); - } - if !content.is_empty() { - blocks.push(AssistantContentBlock::Text { text: content }); - } - blocks.extend(tool_calls.into_iter().enumerate().map( - |(tool_index, tool_call)| { - AssistantContentBlock::ToolCall(AssistantToolCall { - id: tool_call.id.unwrap_or_else(|| { - format!("fixture-tool-call-{index}-{tool_index}") - }), - name: tool_call.function.name, - arguments: tool_call.function.arguments, - }) - }, - )); - ChatMessage::assistant_blocks(blocks) - } - FixtureMessage::Tool { - content, - tool_call_id, - } => ChatMessage::tool_response( - content, - tool_call_id.unwrap_or_else(|| format!("fixture-tool-response-{index}")), - ), - }) - .collect(), - tools: to_chat_tools(&fixture_tools), - tool_choice: if fixture_tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, - ..ChatRequest::for_test() - }; - if matches!( - request.messages.last().map(ChatMessage::role), - Some(ChatRole::Assistant) - ) { - request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; - } - request - .chat_options - .template_kwargs - .insert("thinking".to_string(), Value::Bool(true)); - request + fixture_chat_request(&fixture_path(input_name), deepseek_fixture_options()) } -fn to_chat_tools(tools: &[FixtureTool]) -> Vec { - tools - .iter() - .map(|tool| ChatTool { - name: tool.function.name.clone(), - description: tool.function.description.clone(), - parameters: tool.function.parameters.clone(), - strict: tool.function.strict, - }) - .collect() +fn deepseek_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + enable_thinking: true, + no_generation_prompt_when_last_assistant: true, + } } fn fixture_path(name: &str) -> PathBuf { diff --git a/rust/src/chat/src/renderer/harmony/encoding.rs b/rust/src/chat/src/renderer/harmony/encoding.rs new file mode 100644 index 00000000000..74dfb7cd04a --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/encoding.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Shared Harmony encoding helper for the GPT-OSS renderer and output parser. + +use std::sync::LazyLock; + +use anyhow::Context as _; +use openai_harmony::{HarmonyEncoding, HarmonyEncodingName, load_harmony_encoding}; +use thiserror_ext::AsReport as _; + +use crate::error::{Error, Result}; + +/// Lazily load the shared GPT-OSS Harmony encoding once per process. +pub(crate) fn harmony_encoding() -> Result<&'static HarmonyEncoding> { + static ENCODING: LazyLock> = LazyLock::new(|| { + load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss) + .context("failed to load harmony encoding for gpt-oss") + }); + + ENCODING.as_ref().map_err(|error| Error::HarmonyOutputParsing { + error: error.to_report_string().into(), + }) +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json new file mode 100644 index 00000000000..50edd03ee42 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.json @@ -0,0 +1,14 @@ +{ + "add_generation_prompt": false, + "messages": [ + { + "role": "user", + "content": "What is 2 + 2?" + }, + { + "role": "assistant", + "reasoning_content": "Need simple arithmetic.", + "content": "4" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt new file mode 100644 index 00000000000..dc08897e228 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/assistant_history.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant<|channel|>final<|message|>4<|end|> diff --git a/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json new file mode 100644 index 00000000000..4516e3b32ba --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.json @@ -0,0 +1,27 @@ +[ + { + "role": "developer", + "content": "Use tools when needed.", + "tools": [ + { + "function": { + "name": "lookup", + "description": "Lookup a record.", + "parameters": { + "type": "object", + "properties": { + "id": { + "type": "string" + } + }, + "required": ["id"] + } + } + } + ] + }, + { + "role": "user", + "content": "Find record abc." + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt new file mode 100644 index 00000000000..a63b447f29c --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/developer_tools.txt @@ -0,0 +1,23 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Instructions + +Use tools when needed. + +# Tools + +## functions + +namespace functions { + +// Lookup a record. +type lookup = (_: { +id: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Find record abc.<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json new file mode 100644 index 00000000000..75fd5d8b29e --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.json @@ -0,0 +1,15 @@ +[ + { + "role": "user", + "content": "What is 2 + 2?" + }, + { + "role": "assistant", + "reasoning_content": "This should be dropped.", + "content": "4" + }, + { + "role": "user", + "content": "What is 3 + 5?" + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt new file mode 100644 index 00000000000..9e967b79564 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/drop_analysis.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant<|channel|>final<|message|>4<|end|><|start|>user<|message|>What is 3 + 5?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json new file mode 100644 index 00000000000..5ff190d0b85 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.json @@ -0,0 +1,13 @@ +{ + "reasoning_effort": "high", + "messages": [ + { + "role": "system", + "content": "Answer tersely." + }, + { + "role": "user", + "content": "What is 2 + 2?" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt new file mode 100644 index 00000000000..e656a0a0a47 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/leading_system.txt @@ -0,0 +1,9 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: high + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions + +Answer tersely.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json new file mode 100644 index 00000000000..db5988182fe --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.json @@ -0,0 +1,26 @@ +{ + "tools": [ + { + "function": { + "name": "get_weather", + "description": "Get weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + } + }, + "required": ["city"] + }, + "strict": true + } + } + ], + "messages": [ + { + "role": "user", + "content": "Check Hangzhou weather." + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt new file mode 100644 index 00000000000..f31bc449bfe --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/request_tools.txt @@ -0,0 +1,19 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools + +## functions + +namespace functions { + +// Get weather for a city. +type get_weather = (_: { +city: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Check Hangzhou weather.<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json new file mode 100644 index 00000000000..b8b7f597d6f --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.json @@ -0,0 +1,6 @@ +[ + { + "role": "user", + "content": "Hello, who are you?" + } +] diff --git a/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt new file mode 100644 index 00000000000..7e44ca314ce --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/simple_user.txt @@ -0,0 +1,7 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>Hello, who are you?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt b/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt new file mode 100644 index 00000000000..8ad0ac7d0ee --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/system_instructions_env.txt @@ -0,0 +1,8 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Answer tersely. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: high + +# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>What is 2 + 2?<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json new file mode 100644 index 00000000000..00ccae641e0 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.json @@ -0,0 +1,43 @@ +{ + "tools": [ + { + "function": { + "name": "get_weather", + "description": "Get weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + } + }, + "required": ["city"] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Check Hangzhou weather." + }, + { + "role": "assistant", + "reasoning_content": "Need current weather.", + "tool_calls": [ + { + "id": "call-weather", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call-weather", + "content": "{\"temperature\":20}" + } + ] +} diff --git a/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt new file mode 100644 index 00000000000..0e06a4d107e --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/fixtures/tool_roundtrip.txt @@ -0,0 +1,19 @@ +<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. +Knowledge cutoff: 2024-06 +Current date: 2025-06-28 + +Reasoning: medium + +# Valid channels: analysis, commentary, final. Channel must be included for every message. +Calls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>developer<|message|># Tools + +## functions + +namespace functions { + +// Get weather for a city. +type get_weather = (_: { +city: string, +}) => any; + +} // namespace functions<|end|><|start|>user<|message|>Check Hangzhou weather.<|end|><|start|>assistant<|channel|>analysis<|message|>Need current weather.<|end|><|start|>assistant<|channel|>commentary to=functions.get_weather <|constrain|>json<|message|>{"city":"Hangzhou"}<|call|><|start|>functions.get_weather<|channel|>commentary to=assistant<|message|>{"temperature":20}<|end|><|start|>assistant diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs new file mode 100644 index 00000000000..49375abe8a5 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -0,0 +1,490 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Native Harmony chat renderer for `gpt_oss`. + +pub(crate) mod encoding; + +use openai_harmony::HarmonyEncoding; +use openai_harmony::chat::{ + Author, Conversation, DeveloperContent, Message, ReasoningEffort as HarmonyReasoningEffort, + Role, SystemContent, ToolDescription, +}; +use thiserror_ext::AsReport as _; +use time::macros::format_description; +use vllm_text::Prompt; + +use self::encoding::harmony_encoding; +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::error::{Error, Result}; +use crate::event::AssistantContentBlock; +use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatTool, GenerationPromptMode}; +use crate::{AssistantMessageExt as _, ReasoningEffort}; + +const SYSTEM_START_DATE_ENV: &str = "VLLM_SYSTEM_START_DATE"; +const HARMONY_SYSTEM_INSTRUCTIONS_ENV: &str = "VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS"; + +/// GPT-OSS renderer backed by the official Harmony encoding. +pub struct HarmonyChatRenderer { + encoding: &'static HarmonyEncoding, + options: Options, +} + +struct Options { + system_start_date: String, + use_system_instructions: bool, +} + +impl HarmonyChatRenderer { + /// Create a Harmony renderer for production use. + /// + /// Environment-derived options are resolved once at construction time: + /// + /// - `VLLM_SYSTEM_START_DATE` pins the Harmony system start date. When it is + /// unset, the renderer uses the current local date with a UTC fallback. + /// - `VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS` moves leading instructions + /// into the system model identity when set to a non-zero integer. + pub fn new() -> Result { + Self::with_options( + env_system_start_date(), + env_use_harmony_system_instructions(), + ) + } + + /// Create a Harmony renderer with explicit preamble options. + /// + /// Tests use this constructor to avoid process-global environment mutation. + /// Production code should call [`Self::new`] so the renderer observes the + /// same environment contract as the Python Harmony path. + pub fn with_options( + system_start_date: impl Into, + use_system_instructions: bool, + ) -> Result { + Ok(Self { + encoding: harmony_encoding()?, + options: Options { + system_start_date: system_start_date.into(), + use_system_instructions, + }, + }) + } + + /// Render a chat request directly to Harmony token IDs. + /// + /// Harmony owns both prompt formatting and tokenization, so the Rust + /// frontend bypasses the generic HF tokenizer path for GPT-OSS input. + fn render_token_ids(&self, request: &ChatRequest) -> Result> { + if request.has_multimodal() { + return Err(Error::UnsupportedMultimodalContent("image_url")); + } + if matches!( + request.chat_options.generation_prompt_mode, + GenerationPromptMode::ContinueFinalAssistant + ) { + return Err(Error::ChatTemplate( + "Harmony renderer does not support continue_final_message".to_string(), + )); + } + + let messages = auto_drop_analysis_messages(to_harmony_messages(request, &self.options)?); + let conversation = Conversation::from_messages(messages); + // Pass `None` so oss-harmony does not apply its narrower built-in + // analysis-drop policy after the Rust-side Python-parity cleanup above. + let token_ids = match request.chat_options.generation_prompt_mode { + GenerationPromptMode::StartNewAssistant => self + .encoding + .render_conversation_for_completion(&conversation, Role::Assistant, None), + GenerationPromptMode::NoGenerationPrompt => { + self.encoding.render_conversation(&conversation, None) + } + GenerationPromptMode::ContinueFinalAssistant => unreachable!("checked above"), + } + .map_err(|error| { + Error::ChatTemplate(format!( + "failed to render Harmony prompt: {}", + error.as_report() + )) + })?; + + Ok(token_ids) + } +} + +impl ChatRenderer for HarmonyChatRenderer { + /// Render a chat request as [`Prompt::TokenIds`] with template kwargs echoed + /// for downstream accounting/debugging. + fn render(&self, request: &ChatRequest) -> Result { + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(self.render_token_ids(request)?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} + +/// Convert a vLLM chat request into a full Harmony conversation. +/// +/// This adds the Harmony system/developer preamble, peels at most one leading +/// system/developer instruction message, and then lowers the remaining chat +/// history message-by-message. +fn to_harmony_messages(request: &ChatRequest, options: &Options) -> Result> { + let (instructions, leading_developer_tools, remaining_messages) = + peel_leading_instructions(&request.messages)?; + let tool_call_names = tool_call_names(&request.messages); + let mut messages = + build_harmony_preamble(request, instructions, leading_developer_tools, options)?; + + for message in remaining_messages { + messages.extend(to_harmony_message(message, &tool_call_names, options)?); + } + + Ok(messages) +} + +/// Extract the optional leading instruction message used by the Harmony preamble. +/// +/// Python only peels the first leading `system` or `developer` message. Later +/// system/developer messages stay in the conversation and are lowered normally. +#[allow(clippy::type_complexity)] +fn peel_leading_instructions( + messages: &[ChatMessage], +) -> Result<(Option, Option<&[ChatTool]>, &[ChatMessage])> { + let Some(first) = messages.first() else { + return Ok((None, None, messages)); + }; + + match first { + ChatMessage::System { content } => Ok((Some(flatten_text(content)?), None, &messages[1..])), + ChatMessage::Developer { content, tools } => Ok(( + Some(flatten_text(content)?), + tools.as_deref(), + &messages[1..], + )), + ChatMessage::User { .. } + | ChatMessage::Assistant { .. } + | ChatMessage::ToolResponse { .. } => Ok((None, None, messages)), + } +} + +/// Build the Harmony preamble for one request. +/// +/// The preamble always contains a system message with date and reasoning-effort +/// metadata. Leading instructions live either in the system model identity or in +/// a developer message depending on `use_system_instructions`; request-level and +/// leading developer tools are attached to the developer message. +fn build_harmony_preamble( + request: &ChatRequest, + instructions: Option, + leading_developer_tools: Option<&[ChatTool]>, + options: &Options, +) -> Result> { + let mut messages = vec![Message::from_role_and_content( + Role::System, + system_content( + instructions.as_deref().filter(|_| options.use_system_instructions), + request.chat_options.reasoning_effort, + &options.system_start_date, + )?, + )]; + + let mut developer = DeveloperContent::new(); + let mut has_developer_content = false; + + if !options.use_system_instructions + && let Some(instructions) = instructions.as_deref().filter(|text| !text.is_empty()) + { + developer = developer.with_instructions(instructions); + has_developer_content = true; + } + + let tool_descriptions = preamble_tool_descriptions(request, leading_developer_tools); + if !tool_descriptions.is_empty() { + developer = developer.with_function_tools(tool_descriptions); + has_developer_content = true; + } + + if has_developer_content { + messages.push(Message::from_role_and_content(Role::Developer, developer)); + } + + Ok(messages) +} + +/// Collect request-level and leading developer function tools for the preamble. +fn preamble_tool_descriptions( + request: &ChatRequest, + leading_developer_tools: Option<&[ChatTool]>, +) -> Vec { + let mut tools = Vec::new(); + if request.tool_parsing_enabled() { + tools.extend(to_tool_descriptions(&request.tools)); + } + if let Some(leading_developer_tools) = leading_developer_tools { + tools.extend(to_tool_descriptions(leading_developer_tools)); + } + tools +} + +/// Construct the Harmony system content for the request preamble. +/// +/// Harmony defaults the reasoning effort to `medium` when none is provided, so +/// this only sets an explicit effort after validating vLLM's request value. +fn system_content( + instructions: Option<&str>, + reasoning_effort: Option, + system_start_date: &str, +) -> Result { + let mut content = + SystemContent::new().with_conversation_start_date(system_start_date.to_string()); + + if let Some(reasoning_effort) = reasoning_effort { + content = content.with_reasoning_effort(to_harmony_reasoning_effort(reasoning_effort)?); + } + + if let Some(instructions) = instructions.filter(|text| !text.is_empty()) { + let model_identity = match content.model_identity.as_deref() { + Some(identity) if !identity.is_empty() => format!("{identity}\n{instructions}"), + _ => instructions.to_string(), + }; + content = content.with_model_identity(model_identity); + } + + Ok(content) +} + +/// Lower a single vLLM chat message into one or more Harmony messages. +/// +/// Assistant messages can split into separate analysis, final, commentary, and +/// tool-call messages. Tool responses require the earlier assistant tool-call ID +/// map so the Harmony tool author can include `functions.{name}`. +fn to_harmony_message( + message: &ChatMessage, + tool_call_names: &std::collections::HashMap, + options: &Options, +) -> Result> { + Ok(match message { + ChatMessage::System { content } => { + let instructions = flatten_text(content)?; + vec![system_or_developer_message( + "system", + instructions, + None, + options, + )?] + } + ChatMessage::Developer { content, tools } => { + let instructions = flatten_text(content)?; + vec![developer_message(Some(instructions), tools.as_deref())] + } + ChatMessage::User { content } => { + vec![Message::from_role_and_content( + Role::User, + flatten_text(content)?, + )] + } + ChatMessage::Assistant { content } => assistant_messages(content), + ChatMessage::ToolResponse { + content, + tool_call_id, + } => { + let name = tool_call_names.get(tool_call_id).ok_or_else(|| { + Error::ChatTemplate(format!( + "invalid Harmony tool message: unknown tool_call_id `{tool_call_id}`" + )) + })?; + vec![ + Message::from_author_and_content( + Author::new(Role::Tool, format!("functions.{name}")), + flatten_text(content)?, + ) + .with_channel("commentary") + .with_recipient("assistant"), + ] + } + }) +} + +/// Lower a non-leading system/developer message. +/// +/// Harmony treats most extra system/developer messages as developer +/// instructions. When system-instructions mode is enabled, system messages are +/// rendered as system model-identity additions to match Python. +fn system_or_developer_message( + role: &str, + instructions: String, + tools: Option<&[ChatTool]>, + options: &Options, +) -> Result { + if role == "system" && options.use_system_instructions { + return Ok(Message::from_role_and_content( + Role::System, + system_content(Some(&instructions), None, &options.system_start_date)?, + )); + } + + Ok(developer_message(Some(instructions), tools)) +} + +/// Build a Harmony developer message with optional instructions and function tools. +fn developer_message(instructions: Option, tools: Option<&[ChatTool]>) -> Message { + let mut content = DeveloperContent::new(); + if let Some(instructions) = instructions.filter(|text| !text.is_empty()) { + content = content.with_instructions(instructions); + } + if let Some(tools) = tools { + let tools = to_tool_descriptions(tools); + if !tools.is_empty() { + content = content.with_function_tools(tools); + } + } + Message::from_role_and_content(Role::Developer, content) +} + +/// Lower assistant history into Harmony channels. +/// +/// Plain assistant text goes to `final`. When the assistant has tool calls, +/// visible text goes to `commentary`, reasoning goes to `analysis`, and each +/// function call becomes a `commentary` message to `functions.{name}` with JSON +/// constrained content. +fn assistant_messages(content: &[AssistantContentBlock]) -> Vec { + let mut messages = Vec::new(); + let has_tool_calls = content.has_tool_calls(); + + if has_tool_calls { + let text = content.text(); + if !text.is_empty() { + messages.push( + Message::from_role_and_content(Role::Assistant, text).with_channel("commentary"), + ); + } + } + + if let Some(reasoning) = content.reasoning() { + messages.push( + Message::from_role_and_content(Role::Assistant, reasoning).with_channel("analysis"), + ); + } + + if has_tool_calls { + for tool_call in content.tool_calls() { + messages.push( + Message::from_role_and_content(Role::Assistant, tool_call.arguments.clone()) + .with_channel("commentary") + .with_recipient(format!("functions.{}", tool_call.name)) + .with_content_type("<|constrain|>json"), + ); + } + } else { + let text = content.text(); + if !text.is_empty() { + messages + .push(Message::from_role_and_content(Role::Assistant, text).with_channel("final")); + } + } + + messages +} + +/// Build the tool-call ID to function-name map used by later tool responses. +fn tool_call_names(messages: &[ChatMessage]) -> std::collections::HashMap { + let mut names = std::collections::HashMap::new(); + for message in messages { + let ChatMessage::Assistant { content } = message else { + continue; + }; + for tool_call in content.tool_calls() { + names.insert(tool_call.id.clone(), tool_call.name.clone()); + } + } + names +} + +/// Drop stale assistant analysis messages using vLLM Python's policy. +/// +/// Once an assistant final message exists, earlier analysis messages represent +/// chain-of-thought for completed turns and should not be replayed to the model. +fn auto_drop_analysis_messages(messages: Vec) -> Vec { + // Match vLLM Python's Harmony cleanup: once an assistant final message exists, + // previous assistant analysis messages are stale chain-of-thought and should + // be removed. oss-harmony can also drop analysis with `Some(Default::default())`, + // but that built-in path only triggers when the last assistant message is final + // and drops relative to the first final message, which misses longer multi-turn + // histories with later user/tool turns. + let Some(last_assistant_final_index) = messages.iter().rposition(|message| { + message.author.role == Role::Assistant && message.channel.as_deref() == Some("final") + }) else { + return messages; + }; + + messages + .into_iter() + .enumerate() + .filter_map(|(index, message)| { + (index >= last_assistant_final_index || message.channel.as_deref() != Some("analysis")) + .then_some(message) + }) + .collect() +} + +/// Flatten vLLM text content and reject unsupported multimodal parts. +fn flatten_text(content: &ChatContent) -> Result { + content.try_flatten_to_text() +} + +/// Convert vLLM function tool definitions to Harmony tool descriptions. +fn to_tool_descriptions(tools: &[ChatTool]) -> Vec { + tools + .iter() + .map(|tool| { + ToolDescription::new( + tool.name.clone(), + tool.description.clone().unwrap_or_default(), + Some(tool.parameters.clone()), + ) + }) + .collect() +} + +/// Map supported OpenAI reasoning-effort values onto Harmony's enum. +fn to_harmony_reasoning_effort( + reasoning_effort: ReasoningEffort, +) -> Result { + match reasoning_effort { + ReasoningEffort::Low => Ok(HarmonyReasoningEffort::Low), + ReasoningEffort::Medium => Ok(HarmonyReasoningEffort::Medium), + ReasoningEffort::High => Ok(HarmonyReasoningEffort::High), + ReasoningEffort::None + | ReasoningEffort::Minimal + | ReasoningEffort::XHigh + | ReasoningEffort::Max => Err(Error::ChatTemplate(format!( + "reasoning_effort={:?} is not supported by Harmony. Supported values are: low, medium, high.", + reasoning_effort.as_str() + ))), + } +} + +/// Resolve the system start date from the environment or the current date. +fn env_system_start_date() -> String { + std::env::var(SYSTEM_START_DATE_ENV) + .ok() + .filter(|date| !date.is_empty()) + .unwrap_or_else(current_date) +} + +/// Format today's date as `YYYY-MM-DD`, preferring local time. +fn current_date() -> String { + const DATE_FORMAT: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]"); + let now = time::OffsetDateTime::now_local().unwrap_or_else(|_| time::OffsetDateTime::now_utc()); + now.format(DATE_FORMAT).expect("static date format should be valid") +} + +/// Resolve the env flag that places leading instructions in system identity. +fn env_use_harmony_system_instructions() -> bool { + std::env::var(HARMONY_SYSTEM_INSTRUCTIONS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .is_some_and(|value| value != 0) +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs new file mode 100644 index 00000000000..301d9707e71 --- /dev/null +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::path::PathBuf; + +use expect_test::{ExpectFile, expect, expect_file}; +use thiserror_ext::AsReport as _; + +use super::HarmonyChatRenderer; +use super::encoding::harmony_encoding; +use crate::ChatRenderer; +use crate::error::Error; +use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ + ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, +}; + +const PINNED_DATE: &str = "2025-06-28"; + +fn fixture_request(input_name: &str) -> ChatRequest { + fixture_chat_request( + &fixture_path(input_name), + FixtureRequestOptions { + enable_thinking: false, + no_generation_prompt_when_last_assistant: false, + }, + ) +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/harmony") + .join("fixtures") + .join(name) +} + +fn test_renderer(use_system_instructions: bool) -> HarmonyChatRenderer { + HarmonyChatRenderer::with_options(PINNED_DATE, use_system_instructions).unwrap() +} + +fn render_token_ids(request: &ChatRequest) -> Vec { + render_token_ids_with(&test_renderer(false), request) +} + +fn render_token_ids_with(renderer: &HarmonyChatRenderer, request: &ChatRequest) -> Vec { + renderer + .render(request) + .unwrap() + .prompt + .into_token_ids() + .expect("Harmony renderer returns token IDs") +} + +fn render_prompt_text(request: &ChatRequest) -> String { + render_prompt_text_with(&test_renderer(false), request) +} + +fn render_prompt_text_with(renderer: &HarmonyChatRenderer, request: &ChatRequest) -> String { + let token_ids = render_token_ids_with(renderer, request); + harmony_encoding().unwrap().tokenizer().decode_utf8(&token_ids).unwrap() +} + +fn assert_fixture(input_name: &str, expected: ExpectFile) { + let request = fixture_request(input_name); + let rendered = format!("{}\n", render_prompt_text(&request)); + expected.assert_eq(&rendered); +} + +#[test] +fn renders_token_ids() { + let request = fixture_request("simple_user.json"); + + assert!(!render_token_ids(&request).is_empty()); +} + +#[test] +fn renders_simple_user_fixture() { + assert_fixture("simple_user.json", expect_file!["fixtures/simple_user.txt"]); +} + +#[test] +fn renders_leading_system_fixture() { + assert_fixture( + "leading_system.json", + expect_file!["fixtures/leading_system.txt"], + ); +} + +#[test] +fn renders_system_instructions_env_fixture() { + let renderer = test_renderer(true); + let request = fixture_request("leading_system.json"); + let rendered = format!("{}\n", render_prompt_text_with(&renderer, &request)); + expect_file!["fixtures/system_instructions_env.txt"].assert_eq(&rendered); +} + +#[test] +fn renders_request_tools_fixture() { + assert_fixture( + "request_tools.json", + expect_file!["fixtures/request_tools.txt"], + ); +} + +#[test] +fn renders_developer_tools_fixture() { + assert_fixture( + "developer_tools.json", + expect_file!["fixtures/developer_tools.txt"], + ); +} + +#[test] +fn renders_assistant_history_fixture() { + assert_fixture( + "assistant_history.json", + expect_file!["fixtures/assistant_history.txt"], + ); +} + +#[test] +fn renders_tool_roundtrip_fixture() { + assert_fixture( + "tool_roundtrip.json", + expect_file!["fixtures/tool_roundtrip.txt"], + ); +} + +#[test] +fn drops_stale_analysis_fixture() { + assert_fixture( + "drop_analysis.json", + expect_file!["fixtures/drop_analysis.txt"], + ); +} + +#[test] +fn rejects_invalid_reasoning_effort() { + let mut request = ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect![[r#"chat template error: reasoning_effort="none" is not supported by Harmony. Supported values are: low, medium, high."#]] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn rejects_unknown_tool_response_id() { + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![AssistantContentBlock::ToolCall( + AssistantToolCall { + id: "call-known".to_string(), + name: "lookup".to_string(), + arguments: "{}".to_string(), + }, + )]), + ChatMessage::tool_response("{}", "call-unknown"), + ], + ..ChatRequest::for_test() + }; + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect![ + "chat template error: invalid Harmony tool message: unknown tool_call_id `call-unknown`" + ] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn rejects_multimodal_input() { + let request = ChatRequest { + messages: vec![ChatMessage::user(vec![ChatContentPart::image_url( + "data:image/png;base64,test", + )])], + ..ChatRequest::for_test() + }; + + let error = test_renderer(false).render(&request).unwrap_err(); + + assert!(matches!( + error, + Error::UnsupportedMultimodalContent("image_url") + )); +} + +#[test] +fn rejects_continue_final_assistant() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("write"), + ChatMessage::assistant_text("partial"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let error = test_renderer(false).render(&request).unwrap_err(); + + expect!["chat template error: Harmony renderer does not support continue_final_message"] + .assert_eq(&error.to_report_string()); +} + +#[test] +fn no_generation_prompt_omits_trailing_assistant_start() { + let mut request = fixture_request("simple_user.json"); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_prompt_text(&request); + + assert!(!rendered.ends_with("<|start|>assistant")); +} diff --git a/rust/src/chat/src/renderer/hf/error.rs b/rust/src/chat/src/renderer/hf/error.rs index fcc48c75aba..b173f982392 100644 --- a/rust/src/chat/src/renderer/hf/error.rs +++ b/rust/src/chat/src/renderer/hf/error.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use thiserror::Error as ThisError; #[derive(Debug, ThisError)] diff --git a/rust/src/chat/src/renderer/hf/format.rs b/rust/src/chat/src/renderer/hf/format.rs index 2c990fb37ba..f419b9e3a0c 100644 --- a/rust/src/chat/src/renderer/hf/format.rs +++ b/rust/src/chat/src/renderer/hf/format.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::collections::{HashSet, VecDeque}; use std::fmt; use std::str::FromStr; @@ -361,7 +364,6 @@ mod tests { expect![[r#" template_alpaca.jinja => String - template_baichuan.jinja => String template_chatglm.jinja => String template_chatglm2.jinja => String template_chatml.jinja => String @@ -386,7 +388,6 @@ mod tests { tool_chat_template_llama3.2_pythonic.jinja => String tool_chat_template_llama4_json.jinja => OpenAi tool_chat_template_llama4_pythonic.jinja => OpenAi - tool_chat_template_minimax_m1.jinja => OpenAi tool_chat_template_mistral.jinja => String tool_chat_template_mistral3.jinja => OpenAi tool_chat_template_mistral_parallel.jinja => String diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 47c10c0219e..58ee1ccb7aa 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::collections::HashMap; use serde::Serialize; @@ -14,7 +17,7 @@ use self::format::{ }; use self::template::{CompiledChatTemplate, TemplateContext}; use self::value::{TemplateValue, to_template_value}; -use super::{ChatRenderer, RenderedPrompt}; +use super::{ChatRenderer, RenderedPrompt, effective_template_kwargs}; use crate::error::Result; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; use crate::{ @@ -31,9 +34,14 @@ pub use template::{load_chat_template, resolve_chat_template}; pub use self::format::ChatTemplateContentFormatOption; -#[derive(Debug, Clone)] +/// Template-visible placeholder tokens per supported modality. +/// +/// A `None` token means the loaded model does not support that modality, and +/// content parts of that modality are rejected during rendering. +#[derive(Debug, Clone, Default)] pub struct MultimodalRenderInfo { - pub placeholder_token: String, + pub image_token: Option, + pub video_token: Option, } /// Hugging Face chat-template renderer backed by the local Jinja chat-template @@ -155,11 +163,23 @@ impl HfChatRenderer { effective_template: &CompiledChatTemplate, request: &ChatRequest, ) -> Result { - let messages = to_template_messages( + let mut messages = to_template_messages( &request.messages, effective_template.content_format(), self.multimodal.as_ref(), )?; + + // Handling of `continue_final_message`: + // Append a sentinel tag to the final message content, render as usual, then + // truncate the rendered prompt at the tag so any template suffix after the + // final message content (e.g. the end-of-turn marker) is dropped. + let final_message_text = if request.chat_options.continue_final_message() { + let final_message = messages.last_mut().ok_or(Error::EmptyMessages)?; + Some(append_continue_final_message_tag(final_message)?) + } else { + None + }; + let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools)); trace!( message_count = messages.len(), @@ -169,8 +189,8 @@ impl HfChatRenderer { "applying chat template" ); - let mut merged_template_kwargs = self.default_template_kwargs.clone(); - merged_template_kwargs.extend(request.chat_options.template_kwargs.clone()); + let effective_template_kwargs = + effective_template_kwargs(&self.default_template_kwargs, request); let prompt = effective_template .apply(TemplateContext { messages: &messages, @@ -178,12 +198,18 @@ impl HfChatRenderer { continue_final_message: request.chat_options.continue_final_message(), tools: tools.as_deref(), documents: request.documents.as_deref(), - template_kwargs: Some(&merged_template_kwargs), + template_kwargs: Some(&effective_template_kwargs), special_tokens: self.special_tokens.as_ref(), - reasoning_effort: request.chat_options.reasoning_effort, }) .map_err(|error| Error::ChatTemplate(error.to_report_string()))?; + let prompt = match &final_message_text { + Some(final_message_text) => { + truncate_prompt_at_continue_final_message_tag(prompt, final_message_text)? + } + None => prompt, + }; + trace!( prompt_len = prompt.len(), prompt, "rendered chat template prompt" @@ -191,6 +217,7 @@ impl HfChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + effective_template_kwargs, }) } } @@ -235,6 +262,7 @@ enum TemplateContent { enum TemplateContentPart { Text { text: String }, Image, + Video, } #[derive(Debug, Serialize)] @@ -398,9 +426,17 @@ fn to_template_openai_content( } // All multimodal contents are normalized to `{ "type": }`. ChatContentPart::ImageUrl { .. } => { - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; Ok(TemplateContentPart::Image) } + ChatContentPart::VideoUrl { .. } => { + multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + Ok(TemplateContentPart::Video) + } }) .collect(), } @@ -418,9 +454,16 @@ fn to_template_string_content( match part { ChatContentPart::Text { text } => out.push_str(text), ChatContentPart::ImageUrl { .. } => { - let multimodal = - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; - out.push_str(&multimodal.placeholder_token); + let image_token = multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + out.push_str(image_token); + } + ChatContentPart::VideoUrl { .. } => { + let video_token = multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + out.push_str(video_token); } } } @@ -429,6 +472,74 @@ fn to_template_string_content( } } +/// Sentinel appended to the final message content when `continue_final_message` +/// is requested, used to locate the truncation point in the rendered prompt. +/// +/// Same literal as `transformers`. Occurrences of this string earlier in the +/// prompt are harmless because truncation uses the rightmost match, and the +/// appended sentinel ends up last as long as the template renders messages in +/// order. +const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG "; + +/// Append [`CONTINUE_FINAL_MESSAGE_TAG`] to the trailing text of the final +/// message, returning the original text for post-render validation. +// TODO: transformers v5 also allows continuing a non-`content` field (e.g. +// `reasoning_content`) by passing a field name; only the boolean form is +// supported here. +fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result { + let text = match &mut message.content { + TemplateContent::String(text) => Some(text), + // Pick the last text part in the message. + TemplateContent::OpenAi(parts) => parts.iter_mut().rev().find_map(|part| match part { + TemplateContentPart::Text { text } => Some(text), + TemplateContentPart::Image | TemplateContentPart::Video => None, + }), + }; + let text = text.ok_or_else(|| { + Error::ChatTemplate( + "continue_final_message is set but there is no text to continue \ + in the final message" + .to_string(), + ) + })?; + + let original = text.clone(); + text.push_str(CONTINUE_FINAL_MESSAGE_TAG); + Ok(original) +} + +/// Truncate the rendered prompt at [`CONTINUE_FINAL_MESSAGE_TAG`] so that it +/// ends exactly with the final message content, dropping any template suffix +/// such as end-of-turn markers. +fn truncate_prompt_at_continue_final_message_tag( + mut rendered: String, + final_message_text: &str, +) -> Result { + let tag_loc = rendered + .rfind(CONTINUE_FINAL_MESSAGE_TAG.trim_end()) + .filter(|_| rendered.contains(final_message_text.trim())); + let Some(tag_loc) = tag_loc else { + return Err(Error::ChatTemplate(format!( + "continue_final_message is set but the final message does not appear \ + in the prompt after applying the chat template! This can happen if \ + the chat template deletes portions of the final message. Final \ + message to continue: {}", + final_message_text.trim(), + ))); + }; + + if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) { + // The template preserved spacing, so a plain cut at the tag suffices. + rendered.truncate(tag_loc); + } else { + // The template trimmed the trailing spacing of the message content, so + // apply the same trimming to the retained prefix. + rendered.truncate(tag_loc); + rendered.truncate(rendered.trim_end().len()); + } + Ok(rendered) +} + fn to_template_tools(tools: &[ChatTool]) -> Vec { tools .iter() @@ -490,7 +601,8 @@ mod tests { ) -> Result { HfChatRenderer::new(Some(template.to_string()), HashMap::new(), content_format)? .with_multimodal(Some(MultimodalRenderInfo { - placeholder_token: "".to_string(), + image_token: Some("".to_string()), + video_token: Some("